Skip to content

Commit 9061a1c

Browse files
sreyasreya
committed
prevent cache slowdown during fetches
1 parent 558ad30 commit 9061a1c

File tree

1 file changed

+70
-48
lines changed

1 file changed

+70
-48
lines changed

‎enterprise/wsproxy/keycache.go

Lines changed: 70 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package wsproxy
33
import (
44
"context"
55
"sync"
6+
"sync/atomic"
67
"time"
78

89
"golang.org/x/xerrors"
@@ -18,16 +19,19 @@ import (
1819
var _ cryptokeys.Keycache = &CryptoKeyCache{}
1920

2021
type CryptoKeyCache struct {
21-
ctx context.Context
22-
cancel context.CancelFunc
23-
client *wsproxysdk.Client
24-
logger slog.Logger
25-
Clock quartz.Clock
26-
27-
keysMu sync.RWMutex
28-
keys map[int32]codersdk.CryptoKey
29-
latest codersdk.CryptoKey
30-
closed bool
22+
refreshCtx context.Context
23+
refreshCancel context.CancelFunc
24+
client *wsproxysdk.Client
25+
logger slog.Logger
26+
Clock quartz.Clock
27+
28+
keysMu sync.RWMutex
29+
keys map[int32]codersdk.CryptoKey
30+
latest codersdk.CryptoKey
31+
fetchLock sync.RWMutex
32+
lastFetch time.Time
33+
refresher *quartz.Timer
34+
closed atomic.Bool
3135
}
3236

3337
func NewCryptoKeyCache(ctx context.Context, log slog.Logger, client *wsproxysdk.Client, opts ...func(*CryptoKeyCache)) (*CryptoKeyCache, error) {
@@ -46,21 +50,17 @@ func NewCryptoKeyCache(ctx context.Context, log slog.Logger, client *wsproxysdk.
4650
return nil, xerrors.Errorf("initial fetch: %w", err)
4751
}
4852
cache.keys, cache.latest = m, latest
49-
cache.ctx, cache.cancel = context.WithCancel(ctx)
50-
51-
go cache.refresh()
53+
cache.refresher = cache.Clock.AfterFunc(time.Minute*10, cache.refresh)
5254

5355
return cache, nil
5456
}
5557

5658
func (k *CryptoKeyCache) Signing(ctx context.Context) (codersdk.CryptoKey, error) {
57-
k.keysMu.RLock()
58-
59-
if k.closed {
60-
k.keysMu.RUnlock()
59+
if k.isClosed() {
6160
return codersdk.CryptoKey{}, cryptokeys.ErrClosed
6261
}
6362

63+
k.keysMu.RLock()
6464
latest := k.latest
6565
k.keysMu.RUnlock()
6666

@@ -69,34 +69,31 @@ func (k *CryptoKeyCache) Signing(ctx context.Context) (codersdk.CryptoKey, error
6969
return latest, nil
7070
}
7171

72-
k.keysMu.Lock()
73-
defer k.keysMu.Unlock()
72+
k.fetchLock.Lock()
73+
defer k.fetchLock.Unlock()
7474

75-
if k.closed {
75+
if k.isClosed() {
7676
return codersdk.CryptoKey{}, cryptokeys.ErrClosed
7777
}
7878

79+
k.keysMu.RLock()
7980
if k.latest.CanSign(now) {
81+
k.keysMu.RUnlock()
8082
return k.latest, nil
8183
}
8284

83-
var err error
84-
k.keys, k.latest, err = k.fetch(ctx)
85+
_, latest, err := k.fetch(ctx)
8586
if err != nil {
8687
return codersdk.CryptoKey{}, xerrors.Errorf("fetch: %w", err)
8788
}
8889

89-
if !k.latest.CanSign(now) {
90-
return codersdk.CryptoKey{}, cryptokeys.ErrKeyNotFound
91-
}
92-
93-
return k.latest, nil
90+
return latest, nil
9491
}
9592

9693
func (k *CryptoKeyCache) Verifying(ctx context.Context, sequence int32) (codersdk.CryptoKey, error) {
9794
now := k.Clock.Now()
9895
k.keysMu.RLock()
99-
if k.closed {
96+
if k.isClosed() {
10097
k.keysMu.RUnlock()
10198
return codersdk.CryptoKey{}, cryptokeys.ErrClosed
10299
}
@@ -110,7 +107,7 @@ func (k *CryptoKeyCache) Verifying(ctx context.Context, sequence int32) (codersd
110107
k.keysMu.Lock()
111108
defer k.keysMu.Unlock()
112109

113-
if k.closed {
110+
if k.isClosed() {
114111
return codersdk.CryptoKey{}, cryptokeys.ErrClosed
115112
}
116113

@@ -119,13 +116,12 @@ func (k *CryptoKeyCache) Verifying(ctx context.Context, sequence int32) (codersd
119116
return validKey(key, now)
120117
}
121118

122-
var err error
123-
k.keys, k.latest, err = k.fetch(ctx)
119+
keys, _, err := k.fetch(ctx)
124120
if err != nil {
125121
return codersdk.CryptoKey{}, xerrors.Errorf("fetch: %w", err)
126122
}
127123

128-
key, ok = k.keys[sequence]
124+
key, ok = keys[sequence]
129125
if !ok {
130126
return codersdk.CryptoKey{}, cryptokeys.ErrKeyNotFound
131127
}
@@ -134,28 +130,50 @@ func (k *CryptoKeyCache) Verifying(ctx context.Context, sequence int32) (codersd
134130
}
135131

136132
func (k *CryptoKeyCache) refresh() {
137-
k.Clock.TickerFunc(k.ctx, time.Minute*10, func() error {
138-
kmap, latest, err := k.fetch(k.ctx)
139-
if err != nil {
140-
k.logger.Error(k.ctx, "failed to fetch crypto keys", slog.Error(err))
141-
return nil
142-
}
133+
if k.isClosed() {
134+
return
135+
}
136+
137+
k.keysMu.RLock()
138+
if k.Clock.Now().Sub(k.lastFetch) < time.Minute*10 {
139+
k.keysMu.Unlock()
140+
return
141+
}
142+
143+
k.fetchLock.Lock()
144+
defer k.fetchLock.Unlock()
143145

144-
k.keysMu.Lock()
145-
defer k.keysMu.Unlock()
146-
k.keys = kmap
147-
k.latest = latest
148-
return nil
149-
})
146+
_, _, err := k.fetch(k.refreshCtx)
147+
if err != nil {
148+
k.logger.Error(k.refreshCtx, "fetch crypto keys", slog.Error(err))
149+
return
150+
}
150151
}
151152

152153
func (k *CryptoKeyCache) fetch(ctx context.Context) (map[int32]codersdk.CryptoKey, codersdk.CryptoKey, error) {
154+
153155
keys, err := k.client.CryptoKeys(ctx)
154156
if err != nil {
155157
return nil, codersdk.CryptoKey{}, xerrors.Errorf("get security keys: %w", err)
156158
}
157159

158-
kmap, latest := toKeyMap(keys.CryptoKeys, k.Clock.Now())
160+
if len(keys.CryptoKeys) == 0 {
161+
return nil, codersdk.CryptoKey{}, cryptokeys.ErrKeyNotFound
162+
}
163+
164+
now := k.Clock.Now()
165+
kmap, latest := toKeyMap(keys.CryptoKeys, now)
166+
if !latest.CanSign(now) {
167+
return nil, codersdk.CryptoKey{}, cryptokeys.ErrKeyInvalid
168+
}
169+
170+
k.keysMu.Lock()
171+
defer k.keysMu.Unlock()
172+
173+
k.lastFetch = k.Clock.Now()
174+
k.refresher.Reset(time.Minute * 10)
175+
k.keys, k.latest = kmap, latest
176+
159177
return kmap, latest, nil
160178
}
161179

@@ -179,14 +197,18 @@ func validKey(key codersdk.CryptoKey, now time.Time) (codersdk.CryptoKey, error)
179197
return key, nil
180198
}
181199

200+
func (k *CryptoKeyCache) isClosed() bool {
201+
return k.closed.Load()
202+
}
203+
182204
func (k *CryptoKeyCache) Close() {
183205
k.keysMu.Lock()
184206
defer k.keysMu.Unlock()
185207

186-
if k.closed {
208+
if k.isClosed() {
187209
return
188210
}
189211

190-
k.cancel()
191-
k.closed = true
212+
k.refreshCancel()
213+
k.closed.Store(true)
192214
}

0 commit comments

Comments
 (0)