Skip to content

Commit c88df46

Browse files
kylecarbskylecarbs
committed
Add wsconncache
1 parent b4f9615 commit c88df46

File tree

7 files changed

+307
-16
lines changed

7 files changed

+307
-16
lines changed

‎.vscode/settings.json

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
{
22
"cSpell.words": [
33
"apps",
4+
"awsidentity",
5+
"buildinfo",
46
"buildname",
57
"circbuf",
68
"cliflag",
@@ -17,6 +19,7 @@
1719
"Dsts",
1820
"fatih",
1921
"Formik",
22+
"gitsshkey",
2023
"goarch",
2124
"gographviz",
2225
"goleak",
@@ -31,6 +34,7 @@
3134
"incpatch",
3235
"isatty",
3336
"Jobf",
37+
"Keygen",
3438
"kirsle",
3539
"ldflags",
3640
"manifoldco",
@@ -55,6 +59,7 @@
5559
"retrier",
5660
"rpty",
5761
"sdkproto",
62+
"sdktrace",
5863
"Signup",
5964
"sourcemapped",
6065
"stretchr",
@@ -67,13 +72,16 @@
6772
"tfjson",
6873
"tfstate",
6974
"trimprefix",
75+
"turnconn",
7076
"typegen",
7177
"unconvert",
7278
"Untar",
7379
"VMID",
7480
"weblinks",
7581
"webrtc",
82+
"workspaceagent",
7683
"workspaceapps",
84+
"wsconncache",
7785
"xerrors",
7886
"xstate",
7987
"yamux"

‎coderd/coderd.go

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ import (
2929
"github.com/coder/coder/coderd/rbac"
3030
"github.com/coder/coder/coderd/tracing"
3131
"github.com/coder/coder/coderd/turnconn"
32+
"github.com/coder/coder/coderd/wsconncache"
3233
"github.com/coder/coder/codersdk"
3334
"github.com/coder/coder/site"
3435
)
@@ -80,6 +81,7 @@ func New(options *Options) *API {
8081
Options: options,
8182
Handler: r,
8283
}
84+
api.workspaceAgentCache = wsconncache.New(api.dialWorkspaceAgent, 0)
8385

8486
apiKeyMiddleware := httpmw.ExtractAPIKey(options.Database, &httpmw.OAuth2Configs{
8587
Github: options.GithubOAuth2Config,
@@ -348,18 +350,16 @@ func New(options *Options) *API {
348350
})
349351
})
350352
r.NotFound(site.DefaultHandler().ServeHTTP)
351-
352-
// /workspaceapps/auth
353-
354353
return api
355354
}
356355

357356
type API struct {
358357
*Options
359358

360-
Handler chi.Router
361-
websocketWaitMutex sync.Mutex
362-
websocketWaitGroup sync.WaitGroup
359+
Handler chi.Router
360+
websocketWaitMutex sync.Mutex
361+
websocketWaitGroup sync.WaitGroup
362+
workspaceAgentCache *wsconncache.Cache
363363
}
364364

365365
// Close waits for all WebSocket connections to drain before returning.

‎coderd/workspaceagents.go

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -419,7 +419,9 @@ func (api *API) dialWorkspaceAgent(r *http.Request, agentID uuid.UUID) (*agent.C
419419
if err != nil {
420420
return nil, xerrors.Errorf("negotiate: %w", err)
421421
}
422-
options := &peer.ConnOptions{}
422+
options := &peer.ConnOptions{
423+
Logger: api.Logger.Named("agent-dialer"),
424+
}
423425
options.SettingEngine.SetSrflxAcceptanceMinWait(0)
424426
options.SettingEngine.SetRelayAcceptanceMinWait(0)
425427
// Use the ProxyDialer for the TURN server.

‎coderd/workspaceapps.go

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -113,23 +113,17 @@ func (api *API) workspaceAppsProxyPath(rw http.ResponseWriter, r *http.Request)
113113
return
114114
}
115115

116-
conn, err := api.dialWorkspaceAgent(r, agent.ID)
116+
conn, release, err := api.workspaceAgentCache.Acquire(r, agent.ID)
117117
if err != nil {
118118
httpapi.Write(rw, http.StatusInternalServerError, httpapi.Response{
119119
Message: fmt.Sprintf("dial workspace agent: %s", err),
120120
})
121121
return
122122
}
123+
defer release()
123124

124125
proxy := httputil.NewSingleHostReverseProxy(appURL)
125-
defaultTransport, valid := http.DefaultTransport.(*http.Transport)
126-
if !valid {
127-
panic("dev error: default transport isn't a transport")
128-
}
129-
130-
transport := defaultTransport.Clone()
131-
transport.DialContext = conn.DialContext
132-
proxy.Transport = transport
126+
proxy.Transport = conn.HTTPTransport()
133127
r.URL.Path = chi.URLParam(r, "*")
134128
proxy.ServeHTTP(rw, r)
135129
}

‎coderd/wsconncache/wsconncache.go

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
// Package wsconncache caches workspace agent connections by UUID.
2+
package wsconncache
3+
4+
import (
5+
"context"
6+
"net/http"
7+
"sync"
8+
"time"
9+
10+
"github.com/google/uuid"
11+
"go.uber.org/atomic"
12+
"golang.org/x/xerrors"
13+
14+
"github.com/coder/coder/agent"
15+
)
16+
17+
// New creates a new workspace connection cache that closes
18+
// connections after the inactive timeout provided.
19+
//
20+
// Agent connections are cached due to WebRTC negotiation
21+
// taking a few hundred milliseconds.
22+
func New(dialer Dialer, inactiveTimeout time.Duration) *Cache {
23+
if inactiveTimeout == 0 {
24+
inactiveTimeout = 5 * time.Minute
25+
}
26+
return &Cache{
27+
conns: make(map[uuid.UUID]*Conn),
28+
dialer: dialer,
29+
inactiveTimeout: inactiveTimeout,
30+
}
31+
}
32+
33+
// Dialer creates a new agent connection by ID.
34+
type Dialer func(r *http.Request, id uuid.UUID) (*agent.Conn, error)
35+
36+
// Conn wraps an agent connection with a reusable HTTP transport.
37+
type Conn struct {
38+
*agent.Conn
39+
40+
locks atomic.Uint64
41+
timeoutMutex sync.Mutex
42+
timeout *time.Timer
43+
timeoutCancel context.CancelFunc
44+
transport *http.Transport
45+
}
46+
47+
func (c *Conn) HTTPTransport() *http.Transport {
48+
return c.transport
49+
}
50+
51+
// Close ends the HTTP transport if exists, and closes the agent.
52+
func (c *Conn) Close() error {
53+
if c.transport != nil {
54+
c.transport.CloseIdleConnections()
55+
}
56+
if c.timeout != nil {
57+
c.timeout.Stop()
58+
}
59+
return c.Conn.Close()
60+
}
61+
62+
type Cache struct {
63+
connMutex sync.RWMutex
64+
conns map[uuid.UUID]*Conn
65+
dialer Dialer
66+
inactiveTimeout time.Duration
67+
}
68+
69+
func (c *Cache) Acquire(r *http.Request, id uuid.UUID) (*Conn, func(), error) {
70+
c.connMutex.RLock()
71+
conn, exists := c.conns[id]
72+
c.connMutex.RUnlock()
73+
if !exists {
74+
agentConn, err := c.dialer(r, id)
75+
if err != nil {
76+
return nil, nil, xerrors.Errorf("dial: %w", err)
77+
}
78+
timeoutCtx, timeoutCancelFunc := context.WithCancel(context.Background())
79+
defaultTransport, valid := http.DefaultTransport.(*http.Transport)
80+
if !valid {
81+
panic("dev error: default transport is the wrong type")
82+
}
83+
transport := defaultTransport.Clone()
84+
transport.DialContext = agentConn.DialContext
85+
conn = &Conn{
86+
Conn: agentConn,
87+
timeoutCancel: timeoutCancelFunc,
88+
transport: transport,
89+
}
90+
go func() {
91+
select {
92+
case <-timeoutCtx.Done():
93+
_ = conn.CloseWithError(xerrors.New("cache timeout"))
94+
case <-conn.Closed():
95+
}
96+
c.connMutex.Lock()
97+
delete(c.conns, id)
98+
c.connMutex.Unlock()
99+
}()
100+
c.connMutex.Lock()
101+
c.conns[id] = conn
102+
c.connMutex.Unlock()
103+
}
104+
conn.timeoutMutex.Lock()
105+
defer conn.timeoutMutex.Unlock()
106+
if conn.timeout != nil {
107+
conn.timeout.Stop()
108+
}
109+
conn.locks.Inc()
110+
return conn, func() {
111+
conn.locks.Dec()
112+
conn.timeoutMutex.Lock()
113+
defer conn.timeoutMutex.Unlock()
114+
if conn.timeout != nil {
115+
conn.timeout.Stop()
116+
}
117+
conn.timeout = time.AfterFunc(c.inactiveTimeout, conn.timeoutCancel)
118+
}, nil
119+
}
120+
121+
func (c *Cache) Close() error {
122+
c.connMutex.Lock()
123+
defer c.connMutex.Unlock()
124+
for _, conn := range c.conns {
125+
_ = conn.Close()
126+
}
127+
return nil
128+
}

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