Skip to content

Commit 7798533

Browse files
kylecarbskylecarbs
committed
fix: use a background context when piping derp connections
This was causing boatloads of connects to reestablish every time... See #6746
1 parent eaacc26 commit 7798533

File tree

3 files changed

+17
-11
lines changed

3 files changed

+17
-11
lines changed

‎coderd/workspaceagents.go

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -414,10 +414,8 @@ func (api *API) workspaceAgentListeningPorts(rw http.ResponseWriter, r *http.Req
414414
httpapi.Write(ctx, rw, http.StatusOK, portsResponse)
415415
}
416416

417-
func (api *API) dialWorkspaceAgentTailnet(r *http.Request, agentID uuid.UUID) (*codersdk.WorkspaceAgentConn, error) {
418-
ctx := r.Context()
417+
func (api *API) dialWorkspaceAgentTailnet(ctx context.Context, agentID uuid.UUID) (*codersdk.WorkspaceAgentConn, error) {
419418
clientConn, serverConn := net.Pipe()
420-
421419
conn, err := tailnet.NewConn(&tailnet.Options{
422420
Addresses: []netip.Prefix{netip.PrefixFrom(tailnet.IP(), 128)},
423421
DERPMap: api.DERPMap,
@@ -428,6 +426,7 @@ func (api *API) dialWorkspaceAgentTailnet(r *http.Request, agentID uuid.UUID) (*
428426
_ = serverConn.Close()
429427
return nil, xerrors.Errorf("create tailnet conn: %w", err)
430428
}
429+
ctx, cancel := context.WithCancel(ctx)
431430
conn.SetDERPRegionDialer(func(_ context.Context, region *tailcfg.DERPRegion) net.Conn {
432431
if !region.EmbeddedRelay {
433432
return nil
@@ -437,7 +436,7 @@ func (api *API) dialWorkspaceAgentTailnet(r *http.Request, agentID uuid.UUID) (*
437436
defer left.Close()
438437
defer right.Close()
439438
brw := bufio.NewReadWriter(bufio.NewReader(right), bufio.NewWriter(right))
440-
api.DERPServer.Accept(ctx, right, brw, r.RemoteAddr)
439+
api.DERPServer.Accept(ctx, right, brw, "internal")
441440
}()
442441
return left
443442
})
@@ -453,14 +452,15 @@ func (api *API) dialWorkspaceAgentTailnet(r *http.Request, agentID uuid.UUID) (*
453452
agentConn := &codersdk.WorkspaceAgentConn{
454453
Conn: conn,
455454
CloseFunc: func() {
455+
cancel()
456456
_ = clientConn.Close()
457457
_ = serverConn.Close()
458458
},
459459
}
460460
go func() {
461461
err := (*api.TailnetCoordinator.Load()).ServeClient(serverConn, uuid.New(), agentID)
462462
if err != nil {
463-
api.Logger.Warn(r.Context(), "tailnet coordinator client error", slog.Error(err))
463+
api.Logger.Warn(ctx, "tailnet coordinator client error", slog.Error(err))
464464
_ = agentConn.Close()
465465
}
466466
}()

‎coderd/wsconncache/wsconncache.go

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,15 +24,18 @@ func New(dialer Dialer, inactiveTimeout time.Duration) *Cache {
2424
if inactiveTimeout == 0 {
2525
inactiveTimeout = 5 * time.Minute
2626
}
27+
ctx, cancelFunc := context.WithCancel(context.Background())
2728
return &Cache{
29+
closeContext: ctx,
30+
closeCancel: cancelFunc,
2831
closed: make(chan struct{}),
2932
dialer: dialer,
3033
inactiveTimeout: inactiveTimeout,
3134
}
3235
}
3336

3437
// Dialer creates a new agent connection by ID.
35-
type Dialer func(r *http.Request, id uuid.UUID) (*codersdk.WorkspaceAgentConn, error)
38+
type Dialer func(ctx context.Context, id uuid.UUID) (*codersdk.WorkspaceAgentConn, error)
3639

3740
// Conn wraps an agent connection with a reusable HTTP transport.
3841
type Conn struct {
@@ -66,6 +69,8 @@ type Cache struct {
6669
closed chan struct{}
6770
closeMutex sync.Mutex
6871
closeGroup sync.WaitGroup
72+
closeContext context.Context
73+
closeCancel context.CancelFunc
6974
connGroup singleflight.Group
7075
connMap sync.Map
7176
dialer Dialer
@@ -95,7 +100,7 @@ func (c *Cache) Acquire(r *http.Request, id uuid.UUID) (*Conn, func(), error) {
95100
}
96101
c.closeGroup.Add(1)
97102
c.closeMutex.Unlock()
98-
agentConn, err := c.dialer(r, id)
103+
agentConn, err := c.dialer(c.closeContext, id)
99104
if err != nil {
100105
c.closeGroup.Done()
101106
return nil, xerrors.Errorf("dial: %w", err)
@@ -161,6 +166,7 @@ func (c *Cache) Close() error {
161166
default:
162167
}
163168
close(c.closed)
169+
c.closeCancel()
164170
c.closeGroup.Wait()
165171
return nil
166172
}

‎coderd/wsconncache/wsconncache_test.go

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ func TestCache(t *testing.T) {
4040
t.Parallel()
4141
t.Run("Same", func(t *testing.T) {
4242
t.Parallel()
43-
cache := wsconncache.New(func(r *http.Request, id uuid.UUID) (*codersdk.WorkspaceAgentConn, error) {
43+
cache := wsconncache.New(func(_ context.Context, id uuid.UUID) (*codersdk.WorkspaceAgentConn, error) {
4444
return setupAgent(t, agentsdk.Metadata{}, 0), nil
4545
}, 0)
4646
defer func() {
@@ -55,7 +55,7 @@ func TestCache(t *testing.T) {
5555
t.Run("Expire", func(t *testing.T) {
5656
t.Parallel()
5757
called := atomic.NewInt32(0)
58-
cache := wsconncache.New(func(r *http.Request, id uuid.UUID) (*codersdk.WorkspaceAgentConn, error) {
58+
cache := wsconncache.New(func(_ context.Context, id uuid.UUID) (*codersdk.WorkspaceAgentConn, error) {
5959
called.Add(1)
6060
return setupAgent(t, agentsdk.Metadata{}, 0), nil
6161
}, time.Microsecond)
@@ -74,7 +74,7 @@ func TestCache(t *testing.T) {
7474
})
7575
t.Run("NoExpireWhenLocked", func(t *testing.T) {
7676
t.Parallel()
77-
cache := wsconncache.New(func(r *http.Request, id uuid.UUID) (*codersdk.WorkspaceAgentConn, error) {
77+
cache := wsconncache.New(func(_ context.Context, id uuid.UUID) (*codersdk.WorkspaceAgentConn, error) {
7878
return setupAgent(t, agentsdk.Metadata{}, 0), nil
7979
}, time.Microsecond)
8080
defer func() {
@@ -107,7 +107,7 @@ func TestCache(t *testing.T) {
107107
}()
108108
go server.Serve(random)
109109

110-
cache := wsconncache.New(func(r *http.Request, id uuid.UUID) (*codersdk.WorkspaceAgentConn, error) {
110+
cache := wsconncache.New(func(_ context.Context, id uuid.UUID) (*codersdk.WorkspaceAgentConn, error) {
111111
return setupAgent(t, agentsdk.Metadata{}, 0), nil
112112
}, time.Microsecond)
113113
defer func() {

0 commit comments

Comments
 (0)