Skip to content

Commit c51898a

Browse files
ibetitsmikeibetitsmike
committed
WIP
1 parent 7e143dd commit c51898a

File tree

6 files changed

+213
-105
lines changed

6 files changed

+213
-105
lines changed

‎agent/immortalstreams/backedpipe/backed_pipe.go

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -256,9 +256,20 @@ func (bp *BackedPipe) reconnectLocked() error {
256256
bp.connGen++
257257
bp.state = connected
258258

259+
// Store the generation number before releasing the lock
260+
currentGen := bp.connGen
261+
262+
// Release the lock before calling SetGeneration to avoid deadlock
263+
// SetGeneration acquires its own mutex, and we don't want to hold
264+
// the BackedPipe mutex while waiting for component mutexes
265+
bp.mu.Unlock()
266+
259267
// Update the generation on reader and writer for error reporting
260-
bp.reader.SetGeneration(bp.connGen)
261-
bp.writer.SetGeneration(bp.connGen)
268+
bp.reader.SetGeneration(currentGen)
269+
bp.writer.SetGeneration(currentGen)
270+
271+
// Re-acquire the lock to maintain the function contract
272+
bp.mu.Lock()
262273

263274
return nil
264275
}

‎agent/immortalstreams/handler.go

Lines changed: 24 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package immortalstreams
22

33
import (
44
"context"
5+
"errors"
56
"fmt"
67
"net/http"
78
"strconv"
@@ -74,20 +75,21 @@ func (h *Handler) createStream(w http.ResponseWriter, r *http.Request) {
7475

7576
stream, err := h.manager.CreateStream(ctx, req.TCPPort)
7677
if err != nil {
77-
if strings.Contains(err.Error(), "too many immortal streams") {
78+
switch {
79+
case errors.Is(err, ErrTooManyStreams):
7880
httpapi.Write(ctx, w, http.StatusServiceUnavailable, codersdk.Response{
7981
Message: "Too many Immortal Streams.",
8082
})
8183
return
82-
}
83-
if strings.Contains(err.Error(), "the connection was refused") {
84+
case errors.Is(err, ErrConnRefused):
8485
httpapi.Write(ctx, w, http.StatusNotFound, codersdk.Response{
8586
Message: "The connection was refused.",
8687
})
8788
return
89+
default:
90+
httpapi.InternalServerError(w, err)
91+
return
8892
}
89-
httpapi.InternalServerError(w, err)
90-
return
9193
}
9294

9395
httpapi.Write(ctx, w, http.StatusCreated, stream)
@@ -130,14 +132,16 @@ func (h *Handler) deleteStream(w http.ResponseWriter, r *http.Request) {
130132

131133
err := h.manager.DeleteStream(streamID)
132134
if err != nil {
133-
if strings.Contains(err.Error(), "stream not found") {
135+
switch {
136+
case errors.Is(err, ErrStreamNotFound):
134137
httpapi.Write(ctx, w, http.StatusNotFound, codersdk.Response{
135138
Message: "Stream not found",
136139
})
137140
return
141+
default:
142+
httpapi.InternalServerError(w, err)
143+
return
138144
}
139-
httpapi.InternalServerError(w, err)
140-
return
141145
}
142146

143147
w.WriteHeader(http.StatusNoContent)
@@ -196,9 +200,18 @@ func (h *Handler) handleUpgrade(w http.ResponseWriter, r *http.Request) {
196200
// BackedPipe only needs the reader sequence number for replay
197201
err = h.manager.HandleConnection(streamID, wsConn, readSeqNum)
198202
if err != nil {
199-
h.logger.Error(ctx, "failed to handle connection", slog.Error(err))
200-
conn.Close(websocket.StatusInternalError, err.Error())
201-
return
203+
switch {
204+
case errors.Is(err, ErrStreamNotFound):
205+
conn.Close(websocket.StatusUnsupportedData, "Stream not found")
206+
return
207+
case errors.Is(err, ErrAlreadyConnected):
208+
conn.Close(websocket.StatusPolicyViolation, "Already connected")
209+
return
210+
default:
211+
h.logger.Error(ctx, "failed to handle connection", slog.Error(err))
212+
conn.Close(websocket.StatusInternalError, err.Error())
213+
return
214+
}
202215
}
203216

204217
// Keep the connection open until the context is canceled

‎agent/immortalstreams/manager.go

Lines changed: 30 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,12 @@ package immortalstreams
22

33
import (
44
"context"
5+
"errors"
56
"fmt"
67
"io"
78
"net"
89
"sync"
10+
"syscall"
911
"time"
1012

1113
"github.com/google/uuid"
@@ -16,6 +18,14 @@ import (
1618
"github.com/coder/coder/v2/codersdk"
1719
)
1820

21+
// Package-level sentinel errors
22+
var (
23+
ErrTooManyStreams = xerrors.New("too many streams")
24+
ErrStreamNotFound = xerrors.New("stream not found")
25+
ErrConnRefused = xerrors.New("connection refused")
26+
ErrAlreadyConnected = xerrors.New("already connected")
27+
)
28+
1929
const (
2030
// MaxStreams is the maximum number of immortal streams allowed per agent
2131
MaxStreams = 32
@@ -56,7 +66,7 @@ func (m *Manager) CreateStream(ctx context.Context, port int) (*codersdk.Immorta
5666
// Try to evict a disconnected stream
5767
evicted := m.evictOldestDisconnectedLocked()
5868
if !evicted {
59-
return nil, xerrors.New("too many immortal streams")
69+
return nil, ErrTooManyStreams
6070
}
6171
}
6272

@@ -65,7 +75,7 @@ func (m *Manager) CreateStream(ctx context.Context, port int) (*codersdk.Immorta
6575
conn, err := m.dialer.DialContext(ctx, "tcp", addr)
6676
if err != nil {
6777
if isConnectionRefused(err) {
68-
return nil, xerrors.Errorf("the connection was refused")
78+
return nil, ErrConnRefused
6979
}
7080
return nil, xerrors.Errorf("dial local service: %w", err)
7181
}
@@ -88,13 +98,9 @@ func (m *Manager) CreateStream(ctx context.Context, port int) (*codersdk.Immorta
8898

8999
m.streams[id] = stream
90100

91-
return &codersdk.ImmortalStream{
92-
ID: id,
93-
Name: name,
94-
TCPPort: port,
95-
CreatedAt: stream.createdAt,
96-
LastConnectionAt: stream.createdAt,
97-
}, nil
101+
// Return the API representation of the stream
102+
apiStream := stream.ToAPI()
103+
return &apiStream, nil
98104
}
99105

100106
// GetStream returns a stream by ID
@@ -124,7 +130,7 @@ func (m *Manager) DeleteStream(id uuid.UUID) error {
124130

125131
stream, ok := m.streams[id]
126132
if !ok {
127-
return xerrors.New("stream not found")
133+
return ErrStreamNotFound
128134
}
129135

130136
if err := stream.Close(); err != nil {
@@ -216,17 +222,28 @@ func (m *Manager) HandleConnection(id uuid.UUID, conn io.ReadWriteCloser, readSe
216222
m.mu.RUnlock()
217223

218224
if !ok {
219-
return xerrors.New("stream not found")
225+
return ErrStreamNotFound
220226
}
221227

222228
return stream.HandleReconnect(conn, readSeqNum)
223229
}
224230

225231
// isConnectionRefused checks if an error is a connection refused error
226232
func isConnectionRefused(err error) bool {
233+
// Check for syscall.ECONNREFUSED through error unwrapping
234+
var errno syscall.Errno
235+
if errors.As(err, &errno) && errno == syscall.ECONNREFUSED {
236+
return true
237+
}
238+
239+
// Fallback: check for net.OpError with "dial" operation
227240
var opErr *net.OpError
228-
if xerrors.As(err, &opErr) {
229-
return opErr.Op == "dial"
241+
if errors.As(err, &opErr) && opErr.Op == "dial" {
242+
// Check if the underlying error is ECONNREFUSED
243+
if errors.As(opErr.Err, &errno) && errno == syscall.ECONNREFUSED {
244+
return true
245+
}
230246
}
247+
231248
return false
232249
}

‎agent/immortalstreams/manager_test.go

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package immortalstreams_test
22

33
import (
44
"context"
5+
"errors"
56
"io"
67
"net"
78
"runtime"
@@ -73,7 +74,7 @@ func TestManager_CreateStream(t *testing.T) {
7374
// Use a port that's not listening
7475
_, err := manager.CreateStream(ctx, 65535)
7576
require.Error(t, err)
76-
require.Contains(t, err.Error(), "connection was refused")
77+
require.True(t, errors.Is(err, immortalstreams.ErrConnRefused))
7778
})
7879

7980
t.Run("MaxStreamsLimit", func(t *testing.T) {
@@ -145,7 +146,7 @@ func TestManager_CreateStream(t *testing.T) {
145146
// All streams should be connected, so creating another should fail
146147
_, err = manager.CreateStream(ctx, port)
147148
require.Error(t, err)
148-
require.Contains(t, err.Error(), "too many immortal streams")
149+
require.True(t, errors.Is(err, immortalstreams.ErrTooManyStreams))
149150

150151
// Disconnect one stream
151152
err = manager.DeleteStream(streams[0])
@@ -259,7 +260,7 @@ func TestManager_DeleteStream(t *testing.T) {
259260
// Deleting again should error
260261
err = manager.DeleteStream(stream.ID)
261262
require.Error(t, err)
262-
require.Contains(t, err.Error(), "stream not found")
263+
require.True(t, errors.Is(err, immortalstreams.ErrStreamNotFound))
263264
}
264265

265266
func TestManager_GetStream(t *testing.T) {

‎agent/immortalstreams/stream.go

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -110,8 +110,13 @@ func (r *streamReconnector) Reconnect(ctx context.Context, writerSeqNum uint64)
110110
r.s.reconnectCond.Broadcast()
111111
r.s.mu.Unlock()
112112

113-
// Wait for response from HandleReconnect or context cancellation
113+
// Wait for response from HandleReconnect or context cancellation with timeout
114114
r.s.logger.Info(context.Background(), "reconnect function waiting for response")
115+
116+
// Add a timeout to prevent indefinite hanging
117+
timeout := time.NewTimer(30 * time.Second)
118+
defer timeout.Stop()
119+
115120
select {
116121
case resp := <-responseChan:
117122
r.s.logger.Info(context.Background(), "reconnect function got response",
@@ -129,6 +134,16 @@ func (r *streamReconnector) Reconnect(ctx context.Context, writerSeqNum uint64)
129134
// The stream's Close() method will handle cleanup
130135
r.s.logger.Info(context.Background(), "reconnect function shutdown signal received")
131136
return nil, 0, xerrors.New("stream is shutting down")
137+
case <-timeout.C:
138+
// Timeout occurred - clean up the pending request
139+
r.s.mu.Lock()
140+
if r.s.pendingReconnect != nil {
141+
r.s.pendingReconnect = nil
142+
r.s.handshakePending = false
143+
}
144+
r.s.mu.Unlock()
145+
r.s.logger.Info(context.Background(), "reconnect function timed out")
146+
return nil, 0, xerrors.New("timeout waiting for reconnection response")
132147
}
133148
}
134149

@@ -276,7 +291,7 @@ func (s *Stream) HandleReconnect(clientConn io.ReadWriteCloser, readSeqNum uint6
276291
if s.connected {
277292
s.mu.Unlock()
278293
s.logger.Debug(context.Background(), "another goroutine completed reconnection")
279-
return xerrors.New("stream is already connected")
294+
return ErrAlreadyConnected
280295
}
281296

282297
// Ensure a reconnect attempt is requested while we wait.

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy