This repository was archived by the owner on Aug 30, 2024. It is now read-only.
- Notifications
You must be signed in to change notification settings - Fork18
feat: Add DialCache for key-based connection caching#391
Uh oh!
There was an error while loading.Please reload this page.
Merged
Changes fromall commits
Commits
Show all changes
7 commits Select commitHold shift + click to select a range
91c4f66 feat: Add DialCache for key-based connection caching
kylecarbs857a743 Remove DialOptions
kylecarbs0af7789 Move DialFunc to Dial
kylecarbs6aa8048 Add WS options to dial
kylecarbsb0c0167 Requested changes
kylecarbs1199814 Add comment
kylecarbs844385e Fixup
kylecarbsFile filter
Filter by extension
Conversations
Failed to load comments.
Loading
Uh oh!
There was an error while loading.Please reload this page.
Jump to
Jump to file
Failed to load files.
Loading
Uh oh!
There was an error while loading.Please reload this page.
Diff view
Diff view
There are no files selected for viewing
1 change: 1 addition & 0 deletionsinternal/cmd/tunnel.go
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.Learn more about bidirectional Unicode characters
171 changes: 171 additions & 0 deletionswsnet/cache.go
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,171 @@ | ||
| package wsnet | ||
| import ( | ||
| "context" | ||
| "errors" | ||
| "sync" | ||
| "time" | ||
| "golang.org/x/sync/singleflight" | ||
| ) | ||
| // DialCache constructs a new DialerCache. | ||
| // The cache clears connections that: | ||
| // 1. Are older than the TTL and have no active user-created connections. | ||
| // 2. Have been closed. | ||
| func DialCache(ttl time.Duration) *DialerCache { | ||
| dc := &DialerCache{ | ||
| ttl: ttl, | ||
| closed: make(chan struct{}), | ||
| flightGroup: &singleflight.Group{}, | ||
| mut: sync.RWMutex{}, | ||
| dialers: make(map[string]*Dialer), | ||
| atime: make(map[string]time.Time), | ||
| } | ||
| go dc.init() | ||
| return dc | ||
| } | ||
| type DialerCache struct { | ||
| ttl time.Duration | ||
| flightGroup *singleflight.Group | ||
| closed chan struct{} | ||
| mut sync.RWMutex | ||
| // Key is the "key" of a dialer, which is usually the workspace ID. | ||
| dialers map[string]*Dialer | ||
| atime map[string]time.Time | ||
kylecarbs marked this conversation as resolved. Show resolvedHide resolvedUh oh!There was an error while loading.Please reload this page. | ||
| } | ||
| // init starts the ticker for evicting connections. | ||
| func (d *DialerCache) init() { | ||
| ticker := time.NewTicker(time.Second * 30) | ||
| defer ticker.Stop() | ||
| for { | ||
| select { | ||
| case <-d.closed: | ||
| return | ||
| case <-ticker.C: | ||
| d.evict() | ||
| } | ||
| } | ||
| } | ||
| // evict removes lost/broken/expired connections from the cache. | ||
| func (d *DialerCache) evict() { | ||
| var wg sync.WaitGroup | ||
| d.mut.RLock() | ||
| for key, dialer := range d.dialers { | ||
| wg.Add(1) | ||
| key := key | ||
| dialer := dialer | ||
| go func() { | ||
| defer wg.Done() | ||
| evict := false | ||
| select { | ||
| case <-dialer.Closed(): | ||
| evict = true | ||
| default: | ||
| } | ||
deansheather marked this conversation as resolved. Show resolvedHide resolvedUh oh!There was an error while loading.Please reload this page. | ||
| if dialer.ActiveConnections() == 0 && time.Since(d.atime[key]) >= d.ttl { | ||
| evict = true | ||
| } | ||
| // If we're already evicting there's no point in trying to ping. | ||
| if !evict { | ||
| ctx, cancel := context.WithTimeout(context.Background(), time.Second*15) | ||
| defer cancel() | ||
| err := dialer.Ping(ctx) | ||
| if err != nil { | ||
| evict = true | ||
| } | ||
| } | ||
| if !evict { | ||
| return | ||
| } | ||
| _ = dialer.Close() | ||
| // Ensure after Ping and potential delays that we're still testing against | ||
| // the proper dialer. | ||
| if dialer != d.dialers[key] { | ||
| return | ||
| } | ||
| d.mut.Lock() | ||
| defer d.mut.Unlock() | ||
| delete(d.atime, key) | ||
| delete(d.dialers, key) | ||
| }() | ||
| } | ||
| d.mut.RUnlock() | ||
| wg.Wait() | ||
| } | ||
| // Dial returns a Dialer from the cache if one exists with the key provided, | ||
| // or dials a new connection using the dialerFunc. | ||
kylecarbs marked this conversation as resolved. Show resolvedHide resolvedUh oh!There was an error while loading.Please reload this page. | ||
| // The bool returns whether the connection was found in the cache or not. | ||
| func (d *DialerCache) Dial(ctx context.Context, key string, dialerFunc func() (*Dialer, error)) (*Dialer, bool, error) { | ||
| select { | ||
| case <-d.closed: | ||
| return nil, false, errors.New("cache closed") | ||
| default: | ||
| } | ||
| d.mut.RLock() | ||
| dialer, ok := d.dialers[key] | ||
| d.mut.RUnlock() | ||
| if ok { | ||
| closed := false | ||
| select { | ||
| case <-dialer.Closed(): | ||
| closed = true | ||
| default: | ||
| } | ||
| if !closed { | ||
| d.mut.Lock() | ||
| d.atime[key] = time.Now() | ||
| d.mut.Unlock() | ||
| return dialer, true, nil | ||
| } | ||
| } | ||
| rawDialer, err, _ := d.flightGroup.Do(key, func() (interface{}, error) { | ||
| dialer, err := dialerFunc() | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| d.mut.Lock() | ||
| d.dialers[key] = dialer | ||
| d.atime[key] = time.Now() | ||
| d.mut.Unlock() | ||
| return dialer, nil | ||
| }) | ||
| if err != nil { | ||
| return nil, false, err | ||
| } | ||
| select { | ||
| case <-d.closed: | ||
| return nil, false, errors.New("cache closed") | ||
| default: | ||
| } | ||
| return rawDialer.(*Dialer), false, nil | ||
| } | ||
| // Close closes all cached dialers. | ||
| func (d *DialerCache) Close() error { | ||
| d.mut.Lock() | ||
| defer d.mut.Unlock() | ||
| for _, dialer := range d.dialers { | ||
| err := dialer.Close() | ||
| if err != nil { | ||
| return err | ||
| } | ||
| } | ||
| close(d.closed) | ||
kylecarbs marked this conversation as resolved. Show resolvedHide resolvedUh oh!There was an error while loading.Please reload this page. | ||
| return nil | ||
| } | ||
71 changes: 71 additions & 0 deletionswsnet/cache_test.go
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,71 @@ | ||
| package wsnet | ||
| import ( | ||
| "context" | ||
| "testing" | ||
| "time" | ||
| "cdr.dev/slog/sloggers/slogtest" | ||
| "github.com/stretchr/testify/assert" | ||
| "github.com/stretchr/testify/require" | ||
| ) | ||
| func TestCache(t *testing.T) { | ||
| dialFunc := func(connectAddr string) func() (*Dialer, error) { | ||
| return func() (*Dialer, error) { | ||
| return DialWebsocket(context.Background(), connectAddr, nil, nil) | ||
| } | ||
| } | ||
| t.Run("Caches", func(t *testing.T) { | ||
| connectAddr, listenAddr := createDumbBroker(t) | ||
| l, err := Listen(context.Background(), slogtest.Make(t, nil), listenAddr, "") | ||
| require.NoError(t, err) | ||
| defer l.Close() | ||
| cache := DialCache(time.Hour) | ||
| c1, cached, err := cache.Dial(context.Background(), "example", dialFunc(connectAddr)) | ||
| require.NoError(t, err) | ||
| require.Equal(t, cached, false) | ||
| c2, cached, err := cache.Dial(context.Background(), "example", dialFunc(connectAddr)) | ||
| require.NoError(t, err) | ||
| require.Equal(t, cached, true) | ||
| assert.Same(t, c1, c2) | ||
| }) | ||
| t.Run("Create If Closed", func(t *testing.T) { | ||
| connectAddr, listenAddr := createDumbBroker(t) | ||
| l, err := Listen(context.Background(), slogtest.Make(t, nil), listenAddr, "") | ||
| require.NoError(t, err) | ||
| defer l.Close() | ||
| cache := DialCache(time.Hour) | ||
| c1, cached, err := cache.Dial(context.Background(), "example", dialFunc(connectAddr)) | ||
| require.NoError(t, err) | ||
| require.Equal(t, cached, false) | ||
| require.NoError(t, c1.Close()) | ||
| c2, cached, err := cache.Dial(context.Background(), "example", dialFunc(connectAddr)) | ||
| require.NoError(t, err) | ||
| require.Equal(t, cached, false) | ||
| assert.NotSame(t, c1, c2) | ||
| }) | ||
| t.Run("Evict No Connections", func(t *testing.T) { | ||
| connectAddr, listenAddr := createDumbBroker(t) | ||
| l, err := Listen(context.Background(), slogtest.Make(t, nil), listenAddr, "") | ||
| require.NoError(t, err) | ||
| defer l.Close() | ||
| cache := DialCache(0) | ||
| c1, cached, err := cache.Dial(context.Background(), "example", dialFunc(connectAddr)) | ||
| require.NoError(t, err) | ||
| require.Equal(t, cached, false) | ||
| cache.evict() | ||
| c2, cached, err := cache.Dial(context.Background(), "example", dialFunc(connectAddr)) | ||
| require.NoError(t, err) | ||
| require.Equal(t, cached, false) | ||
kylecarbs marked this conversation as resolved. Show resolvedHide resolvedUh oh!There was an error while loading.Please reload this page. | ||
| assert.NotSame(t, c1, c2) | ||
| }) | ||
| } | ||
12 changes: 9 additions & 3 deletionswsnet/dial.go
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.Learn more about bidirectional Unicode characters
22 changes: 11 additions & 11 deletionswsnet/dial_test.go
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.