Skip to content

Commit ae34934

Browse files
ibetitsmikeibetitsmike
committed
added present workspaces to the tunnel state
1 parent 5dec731 commit ae34934

File tree

4 files changed

+101
-37
lines changed

4 files changed

+101
-37
lines changed

‎tailnet/controllers.go

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -897,6 +897,21 @@ type Workspace struct {
897897
agents map[uuid.UUID]*Agent
898898
}
899899

900+
func (a *Workspace) Clone() Workspace {
901+
agents := make(map[uuid.UUID]*Agent, len(a.agents))
902+
for k, v := range a.agents {
903+
clone := v.Clone()
904+
agents[k] = &clone
905+
}
906+
return Workspace{
907+
ID: a.ID,
908+
Name: a.Name,
909+
Status: a.Status,
910+
ownerUsername: a.ownerUsername,
911+
agents: agents,
912+
}
913+
}
914+
900915
type DNSNameOptions struct {
901916
Suffix string
902917
}

‎tailnet/controllers_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1618,7 +1618,7 @@ func TestTunnelAllWorkspaceUpdatesController_Initial(t *testing.T) {
16181618
cbUpdate := testutil.TryReceive(ctx, t, fUH.ch)
16191619
require.Equal(t, currentState, cbUpdate)
16201620

1621-
// Current recvState should match but shouldn't be a fresh state
1621+
// Current recvState should match
16221622
recvState, err := updateCtrl.CurrentState()
16231623
require.NoError(t, err)
16241624
slices.SortFunc(recvState.UpsertedWorkspaces, func(a, b *tailnet.Workspace) int {

‎vpn/tunnel.go

Lines changed: 23 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@ func NewTunnel(
8888
netLoopDone: make(chan struct{}),
8989
uSendCh: s.sendCh,
9090
agents: map[uuid.UUID]tailnet.Agent{},
91+
workspaces: map[uuid.UUID]tailnet.Workspace{},
9192
clock: quartz.NewReal(),
9293
},
9394
}
@@ -347,7 +348,9 @@ type updater struct {
347348
uSendCh chan<- *TunnelMessage
348349
// agents contains the agents that are currently connected to the tunnel.
349350
agents map[uuid.UUID]tailnet.Agent
350-
conn Conn
351+
// workspaces contains the workspaces to which agents are currently connected via the tunnel.
352+
workspaces map[uuid.UUID]tailnet.Workspace
353+
conn Conn
351354

352355
clock quartz.Clock
353356
}
@@ -399,7 +402,7 @@ func (u *updater) sendUpdateResponse(req *request[*TunnelMessage, *ManagerMessag
399402
func (u *updater) createPeerUpdateLocked(update tailnet.WorkspaceUpdate) *PeerUpdate {
400403
// if the update is a snapshot, we need to process the full state
401404
if update.Kind == tailnet.Snapshot {
402-
processSnapshotUpdate(&update, u.agents)
405+
processSnapshotUpdate(&update, u.agents, u.workspaces)
403406
}
404407

405408
out := &PeerUpdate{
@@ -417,6 +420,12 @@ func (u *updater) createPeerUpdateLocked(update tailnet.WorkspaceUpdate) *PeerUp
417420
for _, agent := range update.DeletedAgents {
418421
delete(u.agents, agent.ID)
419422
}
423+
for _, workspace := range update.UpsertedWorkspaces {
424+
u.workspaces[workspace.ID] = workspace.Clone()
425+
}
426+
for _, workspace := range update.DeletedWorkspaces {
427+
delete(u.workspaces, workspace.ID)
428+
}
420429

421430
for i, ws := range update.UpsertedWorkspaces {
422431
out.UpsertedWorkspaces[i] = &Workspace{
@@ -559,7 +568,7 @@ func (u *updater) netStatusLoop() {
559568
// reconnect to the tailnet API is a full state.
560569
// Without this logic we weren't processing deletes for any workspaces or agents deleted
561570
// while the client was disconnected while the computer was asleep.
562-
func processSnapshotUpdate(update *tailnet.WorkspaceUpdate, agents map[uuid.UUID]tailnet.Agent) {
571+
func processSnapshotUpdate(update *tailnet.WorkspaceUpdate, agents map[uuid.UUID]tailnet.Agent, workspaces map[uuid.UUID]tailnet.Workspace) {
563572
// ignoredWorkspaces is initially populated with the workspaces that are
564573
// in the current update. Later on we populate it with the deleted workspaces too
565574
// so that we don't send duplicate updates. Same applies to ignoredAgents.
@@ -573,23 +582,23 @@ func processSnapshotUpdate(update *tailnet.WorkspaceUpdate, agents map[uuid.UUID
573582
ignoredAgents[agent.ID] = struct{}{}
574583
}
575584
for _, agent := range agents {
576-
if _, ok := ignoredAgents[agent.ID]; !ok {
585+
if _, present := ignoredAgents[agent.ID]; !present {
577586
// delete any current agents that are not in the new update
578587
update.DeletedAgents = append(update.DeletedAgents, &tailnet.Agent{
579588
ID: agent.ID,
580589
Name: agent.Name,
581590
WorkspaceID: agent.WorkspaceID,
582591
})
583-
// if the workspace connected to an agent we're deleting,
584-
// is not present in the fresh state, add it to the deleted workspaces
585-
if _, ok := ignoredWorkspaces[agent.WorkspaceID]; !ok {
586-
update.DeletedWorkspaces = append(update.DeletedWorkspaces, &tailnet.Workspace{
587-
// other fields cannot be populated because the tunnel
588-
// only stores agents and corresponding workspaceIDs
589-
ID: agent.WorkspaceID,
590-
})
591-
ignoredWorkspaces[agent.WorkspaceID] = struct{}{}
592-
}
592+
}
593+
}
594+
for _, workspace := range workspaces {
595+
if _, present := ignoredWorkspaces[workspace.ID]; !present {
596+
update.DeletedWorkspaces = append(update.DeletedWorkspaces, &tailnet.Workspace{
597+
ID: workspace.ID,
598+
Name: workspace.Name,
599+
Status: workspace.Status,
600+
})
601+
ignoredWorkspaces[workspace.ID] = struct{}{}
593602
}
594603
}
595604
}

‎vpn/tunnel_internal_test.go

Lines changed: 62 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -293,6 +293,7 @@ func TestUpdater_createPeerUpdate(t *testing.T) {
293293
ctx: ctx,
294294
netLoopDone: make(chan struct{}),
295295
agents: map[uuid.UUID]tailnet.Agent{},
296+
workspaces: map[uuid.UUID]tailnet.Workspace{},
296297
conn: newFakeConn(tailnet.WorkspaceUpdate{}, hsTime),
297298
}
298299

@@ -727,6 +728,7 @@ func TestProcessFreshState(t *testing.T) {
727728
wsID1 := uuid.New()
728729
wsID2 := uuid.New()
729730
wsID3 := uuid.New()
731+
wsID4 := uuid.New()
730732

731733
agentID1 := uuid.New()
732734
agentID2 := uuid.New()
@@ -738,25 +740,32 @@ func TestProcessFreshState(t *testing.T) {
738740
agent3 := tailnet.Agent{ID: agentID3, Name: "agent3", WorkspaceID: wsID3}
739741
agent4 := tailnet.Agent{ID: agentID4, Name: "agent4", WorkspaceID: wsID1}
740742

741-
ws1 := tailnet.Workspace{ID: wsID1, Name: "ws1"}
742-
ws2 := tailnet.Workspace{ID: wsID2, Name: "ws2"}
743-
ws3 := tailnet.Workspace{ID: wsID3, Name: "ws3"}
743+
ws1 := tailnet.Workspace{ID: wsID1, Name: "ws1", Status: proto.Workspace_RUNNING}
744+
ws2 := tailnet.Workspace{ID: wsID2, Name: "ws2", Status: proto.Workspace_RUNNING}
745+
ws3 := tailnet.Workspace{ID: wsID3, Name: "ws3", Status: proto.Workspace_RUNNING}
746+
ws4 := tailnet.Workspace{ID: wsID4, Name: "ws4", Status: proto.Workspace_RUNNING}
744747

745748
initialAgents := map[uuid.UUID]tailnet.Agent{
746749
agentID1: agent1,
747750
agentID2: agent2,
748751
agentID4: agent4,
749752
}
753+
initialWorkspaces := map[uuid.UUID]tailnet.Workspace{
754+
wsID1: ws1,
755+
wsID2: ws2,
756+
}
750757

751758
tests := []struct {
752-
name string
753-
initialAgents map[uuid.UUID]tailnet.Agent
754-
update *tailnet.WorkspaceUpdate
755-
expectedDelete *tailnet.WorkspaceUpdate // We only care about deletions added by the function
759+
name string
760+
initialAgents map[uuid.UUID]tailnet.Agent
761+
initialWorkspaces map[uuid.UUID]tailnet.Workspace
762+
update *tailnet.WorkspaceUpdate
763+
expectedDelete *tailnet.WorkspaceUpdate // We only care about deletions added by the function
756764
}{
757765
{
758-
name: "NoChange",
759-
initialAgents: initialAgents,
766+
name: "NoChange",
767+
initialAgents: initialAgents,
768+
initialWorkspaces: initialWorkspaces,
760769
update: &tailnet.WorkspaceUpdate{
761770
Kind: tailnet.Snapshot,
762771
UpsertedWorkspaces: []*tailnet.Workspace{&ws1, &ws2},
@@ -770,8 +779,9 @@ func TestProcessFreshState(t *testing.T) {
770779
},
771780
},
772781
{
773-
name: "AgentAdded", // Agent 3 added in update
774-
initialAgents: initialAgents,
782+
name: "AgentAdded", // Agent 3 added in update
783+
initialAgents: initialAgents,
784+
initialWorkspaces: initialWorkspaces,
775785
update: &tailnet.WorkspaceUpdate{
776786
Kind: tailnet.Snapshot,
777787
UpsertedWorkspaces: []*tailnet.Workspace{&ws1, &ws2, &ws3},
@@ -785,8 +795,9 @@ func TestProcessFreshState(t *testing.T) {
785795
},
786796
},
787797
{
788-
name: "AgentRemovedWorkspaceAlsoRemoved", // Agent 2 removed, ws2 also removed
789-
initialAgents: initialAgents,
798+
name: "AgentRemovedWorkspaceAlsoRemoved", // Agent 2 removed, ws2 also removed
799+
initialAgents: initialAgents,
800+
initialWorkspaces: initialWorkspaces,
790801
update: &tailnet.WorkspaceUpdate{
791802
Kind: tailnet.Snapshot,
792803
UpsertedWorkspaces: []*tailnet.Workspace{&ws1}, // ws2 not present
@@ -795,15 +806,18 @@ func TestProcessFreshState(t *testing.T) {
795806
DeletedAgents: []*tailnet.Agent{},
796807
},
797808
expectedDelete: &tailnet.WorkspaceUpdate{
798-
DeletedWorkspaces: []*tailnet.Workspace{{ID: wsID2}}, // Expect ws2 to be deleted
809+
DeletedWorkspaces: []*tailnet.Workspace{
810+
{ID: wsID2, Name: "ws2", Status: proto.Workspace_RUNNING},
811+
}, // Expect ws2 to be deleted
799812
DeletedAgents: []*tailnet.Agent{ // Expect agent2 to be deleted
800813
{ID: agentID2, Name: "agent2", WorkspaceID: wsID2},
801814
},
802815
},
803816
},
804817
{
805-
name: "AgentRemovedWorkspaceStays", // Agent 4 removed, but ws1 stays (due to agent1)
806-
initialAgents: initialAgents,
818+
name: "AgentRemovedWorkspaceStays", // Agent 4 removed, but ws1 stays (due to agent1)
819+
initialAgents: initialAgents,
820+
initialWorkspaces: initialWorkspaces,
807821
update: &tailnet.WorkspaceUpdate{
808822
Kind: tailnet.Snapshot,
809823
UpsertedWorkspaces: []*tailnet.Workspace{&ws1, &ws2}, // ws1 still present
@@ -819,8 +833,9 @@ func TestProcessFreshState(t *testing.T) {
819833
},
820834
},
821835
{
822-
name: "InitialAgentsEmpty",
823-
initialAgents: map[uuid.UUID]tailnet.Agent{}, // Start with no agents known
836+
name: "InitialAgentsEmpty",
837+
initialAgents: map[uuid.UUID]tailnet.Agent{}, // Start with no agents known
838+
initialWorkspaces: map[uuid.UUID]tailnet.Workspace{},
824839
update: &tailnet.WorkspaceUpdate{
825840
Kind: tailnet.Snapshot,
826841
UpsertedWorkspaces: []*tailnet.Workspace{&ws1, &ws2},
@@ -834,8 +849,9 @@ func TestProcessFreshState(t *testing.T) {
834849
},
835850
},
836851
{
837-
name: "UpdateEmpty", // Fresh state says nothing exists
838-
initialAgents: initialAgents,
852+
name: "UpdateEmpty", // Snapshot says nothing exists
853+
initialAgents: initialAgents,
854+
initialWorkspaces: initialWorkspaces,
839855
update: &tailnet.WorkspaceUpdate{
840856
Kind: tailnet.Snapshot,
841857
UpsertedWorkspaces: []*tailnet.Workspace{},
@@ -844,14 +860,35 @@ func TestProcessFreshState(t *testing.T) {
844860
DeletedAgents: []*tailnet.Agent{},
845861
},
846862
expectedDelete: &tailnet.WorkspaceUpdate{ // Expect all initial agents/workspaces to be deleted
847-
DeletedWorkspaces: []*tailnet.Workspace{{ID: wsID1}, {ID: wsID2}}, // ws1 and ws2 deleted
863+
DeletedWorkspaces: []*tailnet.Workspace{
864+
{ID: wsID1, Name: "ws1", Status: proto.Workspace_RUNNING},
865+
{ID: wsID2, Name: "ws2", Status: proto.Workspace_RUNNING},
866+
}, // ws1 and ws2 deleted
848867
DeletedAgents: []*tailnet.Agent{ // agent1, agent2, agent4 deleted
849868
{ID: agentID1, Name: "agent1", WorkspaceID: wsID1},
850869
{ID: agentID2, Name: "agent2", WorkspaceID: wsID2},
851870
{ID: agentID4, Name: "agent4", WorkspaceID: wsID1},
852871
},
853872
},
854873
},
874+
{
875+
name: "WorkspaceWithNoAgents", // Snapshot says nothing exists
876+
initialAgents: initialAgents,
877+
initialWorkspaces: map[uuid.UUID]tailnet.Workspace{wsID1: ws1, wsID2: ws2, wsID4: ws4}, // ws4 has no agents
878+
update: &tailnet.WorkspaceUpdate{
879+
Kind: tailnet.Snapshot,
880+
UpsertedWorkspaces: []*tailnet.Workspace{&ws1, &ws2},
881+
UpsertedAgents: []*tailnet.Agent{&agent1, &agent2, &agent4},
882+
DeletedWorkspaces: []*tailnet.Workspace{},
883+
DeletedAgents: []*tailnet.Agent{},
884+
},
885+
expectedDelete: &tailnet.WorkspaceUpdate{ // Expect all initial agents/workspaces to be deleted
886+
DeletedWorkspaces: []*tailnet.Workspace{
887+
{ID: wsID4, Name: "ws4", Status: proto.Workspace_RUNNING},
888+
}, // ws4 should be deleted
889+
DeletedAgents: []*tailnet.Agent{},
890+
},
891+
},
855892
}
856893

857894
for _, tt := range tests {
@@ -862,7 +899,10 @@ func TestProcessFreshState(t *testing.T) {
862899
agentsCopy := make(map[uuid.UUID]tailnet.Agent)
863900
maps.Copy(agentsCopy, tt.initialAgents)
864901

865-
processSnapshotUpdate(tt.update, agentsCopy)
902+
workspaceCopy := make(map[uuid.UUID]tailnet.Workspace)
903+
maps.Copy(workspaceCopy, tt.initialWorkspaces)
904+
905+
processSnapshotUpdate(tt.update, agentsCopy, workspaceCopy)
866906

867907
require.ElementsMatch(t, tt.expectedDelete.DeletedAgents, tt.update.DeletedAgents, "DeletedAgents mismatch")
868908
require.ElementsMatch(t, tt.expectedDelete.DeletedWorkspaces, tt.update.DeletedWorkspaces, "DeletedWorkspaces mismatch")

0 commit comments

Comments
 (0)