Skip to content

Commit 9faa940

Browse files
ethanndicksonethanndickson
committed
use peer ids as subscription key + tests
1 parent 8b55da8 commit 9faa940

File tree

4 files changed

+358
-215
lines changed

4 files changed

+358
-215
lines changed

‎tailnet/service.go

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -259,10 +259,12 @@ func (s *DRPCService) WorkspaceUpdates(_ *proto.WorkspaceUpdatesRequest, stream
259259
)
260260
switch auth := streamID.Auth.(type) {
261261
case ClientUserCoordinateeAuth:
262-
updatesCh, err = s.WorkspaceUpdatesProvider.Subscribe(auth.UserID)
262+
// Stream ID is the peer ID
263+
updatesCh, err = s.WorkspaceUpdatesProvider.Subscribe(streamID.ID, auth.UserID)
263264
if err != nil {
264265
err = xerrors.Errorf("subscribe to workspace updates: %w", err)
265266
}
267+
defer s.WorkspaceUpdatesProvider.Unsubscribe(streamID.ID)
266268
default:
267269
err = xerrors.Errorf("workspace updates not supported by auth name %T", auth)
268270
}
@@ -273,6 +275,9 @@ func (s *DRPCService) WorkspaceUpdates(_ *proto.WorkspaceUpdatesRequest, stream
273275
for {
274276
select {
275277
case updates := <-updatesCh:
278+
if updates == nil {
279+
return nil
280+
}
276281
err := stream.Send(updates)
277282
if err != nil {
278283
return xerrors.Errorf("send workspace update: %w", err)

‎tailnet/workspaceupdates.go

Lines changed: 53 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ func convertRows(v []database.GetWorkspacesAndAgentsRow) workspacesByOwner {
4040
WorkspaceName: ws.Name,
4141
JobStatus: ws.JobStatus,
4242
Transition: ws.Transition,
43+
Agents: ws.Agents,
4344
}
4445
if byID, exists := m[ws.OwnerID]; !exists {
4546
m[ws.OwnerID] = map[uuid.UUID]ownedWorkspace{ws.ID: owned}
@@ -68,14 +69,16 @@ func (s *sub) send(all workspacesByOwner) {
6869
defer s.mu.Unlock()
6970

7071
// Filter to only the workspaces owned by the user
71-
own := all[s.userID]
72-
update := produceUpdate(s.prev, own)
73-
s.prev = own
72+
latest := all[s.userID]
73+
update := produceUpdate(s.prev, latest)
74+
s.prev = latest
7475
s.tx <- update
7576
}
7677

7778
type WorkspaceUpdatesProvider interface {
78-
Subscribe(userID uuid.UUID) (<-chan *proto.WorkspaceUpdate, error)
79+
Subscribe(peerID uuid.UUID, userID uuid.UUID) (<-chan *proto.WorkspaceUpdate, error)
80+
Unsubscribe(peerID uuid.UUID)
81+
Stop()
7982
}
8083

8184
type WorkspaceStore interface {
@@ -84,25 +87,40 @@ type WorkspaceStore interface {
8487
}
8588

8689
type updatesProvider struct {
87-
mu sync.RWMutex
88-
db WorkspaceStore
89-
ps pubsub.Pubsub
90-
subs []*sub
90+
mu sync.RWMutex
91+
db WorkspaceStore
92+
ps pubsub.Pubsub
93+
// Peer ID -> subscription
94+
subs map[uuid.UUID]*sub
95+
latest workspacesByOwner
9196
cancelFn func()
9297
}
9398

9499
var _ WorkspaceUpdatesProvider = (*updatesProvider)(nil)
95100

96-
func (u *updatesProvider) Start() error {
97-
cancel, err := u.ps.Subscribe(codersdk.AllWorkspacesNotifyChannel, u.handleUpdate)
101+
func NewUpdatesProvider(ctx context.Context, db WorkspaceStore, ps pubsub.Pubsub) (WorkspaceUpdatesProvider, error) {
102+
rows, err := db.GetWorkspacesAndAgents(ctx)
103+
if err != nil && !xerrors.Is(err, sql.ErrNoRows) {
104+
return nil, err
105+
}
106+
out := &updatesProvider{
107+
db: db,
108+
ps: ps,
109+
subs: map[uuid.UUID]*sub{},
110+
latest: convertRows(rows),
111+
}
112+
cancel, err := ps.Subscribe(codersdk.AllWorkspacesNotifyChannel, out.handleUpdate)
98113
if err != nil {
99-
return err
114+
return nil, err
100115
}
101-
u.cancelFn = cancel
102-
return nil
116+
out.cancelFn = cancel
117+
return out, nil
103118
}
104119

105120
func (u *updatesProvider) Stop() {
121+
for _, sub := range u.subs {
122+
close(sub.tx)
123+
}
106124
u.cancelFn()
107125
}
108126

@@ -116,7 +134,6 @@ func (u *updatesProvider) handleUpdate(ctx context.Context, _ []byte) {
116134
wg := &sync.WaitGroup{}
117135
latest := convertRows(rows)
118136
u.mu.RLock()
119-
defer u.mu.RUnlock()
120137
for _, sub := range u.subs {
121138
sub := sub
122139
wg.Add(1)
@@ -125,18 +142,15 @@ func (u *updatesProvider) handleUpdate(ctx context.Context, _ []byte) {
125142
defer wg.Done()
126143
}()
127144
}
128-
wg.Wait()
129-
}
145+
u.mu.RUnlock()
130146

131-
func NewUpdatesProvider(db WorkspaceStore, ps pubsub.Pubsub) WorkspaceUpdatesProvider {
132-
return &updatesProvider{
133-
db: db,
134-
ps: ps,
135-
subs: make([]*sub, 0),
136-
}
147+
u.mu.Lock()
148+
u.latest = latest
149+
u.mu.Unlock()
150+
wg.Wait()
137151
}
138152

139-
func (u *updatesProvider) Subscribe(userID uuid.UUID) (<-chan *proto.WorkspaceUpdate, error) {
153+
func (u *updatesProvider) Subscribe(peerID uuid.UUID, userID uuid.UUID) (<-chan *proto.WorkspaceUpdate, error) {
140154
u.mu.Lock()
141155
defer u.mu.Unlock()
142156

@@ -146,10 +160,25 @@ func (u *updatesProvider) Subscribe(userID uuid.UUID) (<-chan *proto.WorkspaceUp
146160
tx: tx,
147161
prev: make(workspacesByID),
148162
}
149-
u.subs = append(u.subs, sub)
163+
u.subs[peerID] = sub
164+
// Write initial state
165+
sub.send(u.latest)
150166
return tx, nil
151167
}
152168

169+
func (u *updatesProvider) Unsubscribe(peerID uuid.UUID) {
170+
u.mu.Lock()
171+
defer u.mu.Unlock()
172+
173+
sub, exists := u.subs[peerID]
174+
if !exists {
175+
return
176+
}
177+
close(sub.tx)
178+
delete(u.subs, peerID)
179+
return
180+
}
181+
153182
func produceUpdate(old, new workspacesByID) *proto.WorkspaceUpdate {
154183
out := &proto.WorkspaceUpdate{
155184
UpsertedWorkspaces: []*proto.Workspace{},

‎tailnet/workspaceupdates_internal_test.go

Lines changed: 0 additions & 190 deletions
This file was deleted.

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