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

Skip to content

Commit e95fa44

Browse files
ibetitsmikeibetitsmike
committed
initial implementation
1 parent b47d54d commit e95fa44

File tree

6 files changed

+208
-72
lines changed

6 files changed

+208
-72
lines changed

‎tailnet/controllers.go

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1049,6 +1049,7 @@ func (t *tunnelUpdater) recvLoop() {
10491049
t.logger.Debug(context.Background(), "tunnel updater recvLoop started")
10501050
defer t.logger.Debug(context.Background(), "tunnel updater recvLoop done")
10511051
defer close(t.recvLoopDone)
1052+
freshState := true
10521053
for {
10531054
update, err := t.client.Recv()
10541055
if err != nil {
@@ -1061,8 +1062,10 @@ func (t *tunnelUpdater) recvLoop() {
10611062
}
10621063
t.logger.Debug(context.Background(), "got workspace update",
10631064
slog.F("workspace_update", update),
1065+
slog.F("fresh_state", freshState),
10641066
)
1065-
err = t.handleUpdate(update)
1067+
err = t.handleUpdate(update, freshState)
1068+
freshState = false
10661069
if err != nil {
10671070
t.logger.Critical(context.Background(), "failed to handle workspace Update", slog.Error(err))
10681071
cErr := t.client.Close()
@@ -1083,6 +1086,7 @@ type WorkspaceUpdate struct {
10831086
UpsertedAgents []*Agent
10841087
DeletedWorkspaces []*Workspace
10851088
DeletedAgents []*Agent
1089+
FreshState bool
10861090
}
10871091

10881092
func (w *WorkspaceUpdate) Clone() WorkspaceUpdate {
@@ -1091,6 +1095,7 @@ func (w *WorkspaceUpdate) Clone() WorkspaceUpdate {
10911095
UpsertedAgents: make([]*Agent, len(w.UpsertedAgents)),
10921096
DeletedWorkspaces: make([]*Workspace, len(w.DeletedWorkspaces)),
10931097
DeletedAgents: make([]*Agent, len(w.DeletedAgents)),
1098+
FreshState: w.FreshState,
10941099
}
10951100
for i, ws := range w.UpsertedWorkspaces {
10961101
clone.UpsertedWorkspaces[i] = &Workspace{
@@ -1115,7 +1120,7 @@ func (w *WorkspaceUpdate) Clone() WorkspaceUpdate {
11151120
return clone
11161121
}
11171122

1118-
func (t *tunnelUpdater) handleUpdate(update *proto.WorkspaceUpdate) error {
1123+
func (t *tunnelUpdater) handleUpdate(update *proto.WorkspaceUpdate, freshState bool) error {
11191124
t.Lock()
11201125
defer t.Unlock()
11211126

@@ -1124,6 +1129,7 @@ func (t *tunnelUpdater) handleUpdate(update *proto.WorkspaceUpdate) error {
11241129
UpsertedAgents: []*Agent{},
11251130
DeletedWorkspaces: []*Workspace{},
11261131
DeletedAgents: []*Agent{},
1132+
FreshState: freshState,
11271133
}
11281134

11291135
for _, uw := range update.UpsertedWorkspaces {

‎tailnet/controllers_test.go

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1611,13 +1611,15 @@ func TestTunnelAllWorkspaceUpdatesController_Initial(t *testing.T) {
16111611
},
16121612
DeletedWorkspaces: []*tailnet.Workspace{},
16131613
DeletedAgents: []*tailnet.Agent{},
1614+
FreshState: true,
16141615
}
16151616

16161617
// And the callback
16171618
cbUpdate := testutil.TryReceive(ctx, t, fUH.ch)
16181619
require.Equal(t, currentState, cbUpdate)
16191620

1620-
// Current recvState should match
1621+
// Current recvState should match but shouldn't be a fresh state
1622+
currentState.FreshState = false
16211623
recvState, err := updateCtrl.CurrentState()
16221624
require.NoError(t, err)
16231625
slices.SortFunc(recvState.UpsertedWorkspaces, func(a, b *tailnet.Workspace) int {
@@ -1692,12 +1694,14 @@ func TestTunnelAllWorkspaceUpdatesController_DeleteAgent(t *testing.T) {
16921694
},
16931695
DeletedWorkspaces: []*tailnet.Workspace{},
16941696
DeletedAgents: []*tailnet.Agent{},
1697+
FreshState: true,
16951698
}
16961699

16971700
cbUpdate := testutil.TryReceive(ctx, t, fUH.ch)
16981701
require.Equal(t, initRecvUp, cbUpdate)
16991702

1700-
// Current state should match initial
1703+
// Current state should match initial but shouldn't be a fresh state
1704+
initRecvUp.FreshState = false
17011705
state, err := updateCtrl.CurrentState()
17021706
require.NoError(t, err)
17031707
require.Equal(t, initRecvUp, state)
@@ -1753,6 +1757,7 @@ func TestTunnelAllWorkspaceUpdatesController_DeleteAgent(t *testing.T) {
17531757
"w1.coder.": {ws1a1IP},
17541758
}},
17551759
},
1760+
FreshState: false,
17561761
}
17571762
require.Equal(t, sndRecvUpdate, cbUpdate)
17581763

@@ -1771,6 +1776,7 @@ func TestTunnelAllWorkspaceUpdatesController_DeleteAgent(t *testing.T) {
17711776
},
17721777
DeletedWorkspaces: []*tailnet.Workspace{},
17731778
DeletedAgents: []*tailnet.Agent{},
1779+
FreshState: false,
17741780
}, state)
17751781
}
17761782

‎tailnet/proto/tailnet.pb.go

Lines changed: 67 additions & 56 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎tailnet/proto/tailnet.proto

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -207,6 +207,7 @@ message WorkspaceUpdate {
207207
repeated Agent upserted_agents = 2;
208208
repeated Workspace deleted_workspaces = 3;
209209
repeated Agent deleted_agents = 4;
210+
bool fresh_state = 5;
210211
}
211212

212213
message Workspace {

‎vpn/tunnel.go

Lines changed: 19 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -404,7 +404,24 @@ func (u *updater) createPeerUpdateLocked(update tailnet.WorkspaceUpdate) *PeerUp
404404
DeletedAgents: make([]*Agent, len(update.DeletedAgents)),
405405
}
406406

407-
u.saveUpdateLocked(update)
407+
// save the workspace update to the tunnel's state, such that it can
408+
// be used to populate automated peer updates.
409+
for _, agent := range update.UpsertedAgents {
410+
u.agents[agent.ID] = agent.Clone()
411+
}
412+
for _, agent := range update.DeletedAgents {
413+
delete(u.agents, agent.ID)
414+
}
415+
416+
// this flag is true on the first update after a reconnect
417+
if update.FreshState {
418+
for _, agent := range update.UpsertedAgents {
419+
// delete any current agents that are not in the new update
420+
if _, ok := u.agents[agent.ID]; !ok {
421+
delete(u.agents, agent.ID)
422+
}
423+
}
424+
}
408425

409426
for i, ws := range update.UpsertedWorkspaces {
410427
out.UpsertedWorkspaces[i] = &Workspace{
@@ -413,6 +430,7 @@ func (u *updater) createPeerUpdateLocked(update tailnet.WorkspaceUpdate) *PeerUp
413430
Status: Workspace_Status(ws.Status),
414431
}
415432
}
433+
416434
upsertedAgents := u.convertAgentsLocked(update.UpsertedAgents)
417435
out.UpsertedAgents = upsertedAgents
418436
for i, ws := range update.DeletedWorkspaces {
@@ -472,17 +490,6 @@ func (u *updater) convertAgentsLocked(agents []*tailnet.Agent) []*Agent {
472490
return out
473491
}
474492

475-
// saveUpdateLocked saves the workspace update to the tunnel's state, such that it can
476-
// be used to populate automated peer updates.
477-
func (u *updater) saveUpdateLocked(update tailnet.WorkspaceUpdate) {
478-
for _, agent := range update.UpsertedAgents {
479-
u.agents[agent.ID] = agent.Clone()
480-
}
481-
for _, agent := range update.DeletedAgents {
482-
delete(u.agents, agent.ID)
483-
}
484-
}
485-
486493
// setConn sets the `conn` and returns false if there's already a connection set.
487494
func (u *updater) setConn(conn Conn) bool {
488495
u.mu.Lock()

0 commit comments

Comments
 (0)