Thanks to visit codestin.com
Credit goes to github.com

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)