Skip to content

feat: improve coder connect tunnel handling on reconnect #17598

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 10 commits into from
May 6, 2025
Merged

feat: improve coder connect tunnel handling on reconnect #17598

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
52f1c2b
initial implementation
ibetitsmike Apr 29, 2025
288c33e
moved out the logic to separate function
ibetitsmike Apr 29, 2025
0de5df2
additional tests
ibetitsmike Apr 29, 2025
2b77b8e
fixed map copy
ibetitsmike Apr 29, 2025
3d39713
fmt & lint fixes
ibetitsmike Apr 29, 2025
ca2e1bf
removed unecessary test for verifying existing deletetions (fresh sta…
ibetitsmike Apr 30, 2025
5dec731
refactored FreshState to UpdateKind to be a better indicator of the f…
ibetitsmike May 5, 2025
ae34934
added present workspaces to the tunnel state
ibetitsmike May 5, 2025
101702d
fix lint; wrong receiver name
ibetitsmike May 5, 2025
0d2ecb2
compile error mishap
ibetitsmike May 5, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
added present workspaces to the tunnel state
  • Loading branch information
@ibetitsmike
ibetitsmike committed May 5, 2025
commit ae34934ac6ce8de7d66bb11f1d64ed684ea539df
15 changes: 15 additions & 0 deletions tailnet/controllers.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -897,6 +897,21 @@
agents map[uuid.UUID]*Agent
}

func (a *Workspace) Clone() Workspace {
agents := make(map[uuid.UUID]*Agent, len(a.agents))
for k, v := range a.agents {
clone := v.Clone()
agents[k] = &clone
}
return Workspace{
ID: a.ID,
Name: a.Name,
Status: a.Status,
ownerUsername: a.ownerUsername,
agents: agents,
}
}

type DNSNameOptions struct {
Suffix string
}
Expand All @@ -906,7 +921,7 @@
// Usernames are globally unique & case-insensitive.
// Workspace names are unique per-user & case-insensitive.
// Agent names are unique per-workspace & case-insensitive.
func (w *Workspace) updateDNSNames(options DNSNameOptions) error {

Check failure on line 924 in tailnet/controllers.go

View workflow job for this annotation

GitHub Actions / lint 

receiver-naming: receiver name w should be consistent with previous receiver name a for Workspace (revive)
wsName := strings.ToLower(w.Name)
username := strings.ToLower(w.ownerUsername)
for id, a := range w.agents {
Expand Down
2 changes: 1 addition & 1 deletion tailnet/controllers_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1618,7 +1618,7 @@ func TestTunnelAllWorkspaceUpdatesController_Initial(t *testing.T) {
cbUpdate := testutil.TryReceive(ctx, t, fUH.ch)
require.Equal(t, currentState, cbUpdate)

// Current recvState should match but shouldn't be a fresh state
// Current recvState should match
recvState, err := updateCtrl.CurrentState()
require.NoError(t, err)
slices.SortFunc(recvState.UpsertedWorkspaces, func(a, b *tailnet.Workspace) int {
Expand Down
37 changes: 23 additions & 14 deletions vpn/tunnel.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,7 @@ func NewTunnel(
netLoopDone: make(chan struct{}),
uSendCh: s.sendCh,
agents: map[uuid.UUID]tailnet.Agent{},
workspaces: map[uuid.UUID]tailnet.Workspace{},
clock: quartz.NewReal(),
},
}
Expand Down Expand Up @@ -347,7 +348,9 @@ type updater struct {
uSendCh chan<- *TunnelMessage
// agents contains the agents that are currently connected to the tunnel.
agents map[uuid.UUID]tailnet.Agent
conn Conn
// workspaces contains the workspaces to which agents are currently connected via the tunnel.
workspaces map[uuid.UUID]tailnet.Workspace
conn Conn

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

out := &PeerUpdate{
Expand All @@ -417,6 +420,12 @@ func (u *updater) createPeerUpdateLocked(update tailnet.WorkspaceUpdate) *PeerUp
for _, agent := range update.DeletedAgents {
delete(u.agents, agent.ID)
}
for _, workspace := range update.UpsertedWorkspaces {
u.workspaces[workspace.ID] = workspace.Clone()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't actually need a full clone of the workspace, but it's probably clearer to do the legit Clone so we don't have a partially cloned object waiting to trip up some future editor of this code.

}
for _, workspace := range update.DeletedWorkspaces {
delete(u.workspaces, workspace.ID)
}

for i, ws := range update.UpsertedWorkspaces {
out.UpsertedWorkspaces[i] = &Workspace{
Expand Down Expand Up @@ -559,7 +568,7 @@ func (u *updater) netStatusLoop() {
// reconnect to the tailnet API is a full state.
// Without this logic we weren't processing deletes for any workspaces or agents deleted
// while the client was disconnected while the computer was asleep.
func processSnapshotUpdate(update *tailnet.WorkspaceUpdate, agents map[uuid.UUID]tailnet.Agent) {
func processSnapshotUpdate(update *tailnet.WorkspaceUpdate, agents map[uuid.UUID]tailnet.Agent, workspaces map[uuid.UUID]tailnet.Workspace) {
// ignoredWorkspaces is initially populated with the workspaces that are
// in the current update. Later on we populate it with the deleted workspaces too
// so that we don't send duplicate updates. Same applies to ignoredAgents.
Expand All @@ -573,23 +582,23 @@ func processSnapshotUpdate(update *tailnet.WorkspaceUpdate, agents map[uuid.UUID
ignoredAgents[agent.ID] = struct{}{}
}
for _, agent := range agents {
if _, ok := ignoredAgents[agent.ID]; !ok {
if _, present := ignoredAgents[agent.ID]; !present {
// delete any current agents that are not in the new update
update.DeletedAgents = append(update.DeletedAgents, &tailnet.Agent{
ID: agent.ID,
Name: agent.Name,
WorkspaceID: agent.WorkspaceID,
})
// if the workspace connected to an agent we're deleting,
// is not present in the fresh state, add it to the deleted workspaces
if _, ok := ignoredWorkspaces[agent.WorkspaceID]; !ok {
update.DeletedWorkspaces = append(update.DeletedWorkspaces, &tailnet.Workspace{
// other fields cannot be populated because the tunnel
// only stores agents and corresponding workspaceIDs
ID: agent.WorkspaceID,
})
ignoredWorkspaces[agent.WorkspaceID] = struct{}{}
}
}
}
for _, workspace := range workspaces {
if _, present := ignoredWorkspaces[workspace.ID]; !present {
update.DeletedWorkspaces = append(update.DeletedWorkspaces, &tailnet.Workspace{
ID: workspace.ID,
Name: workspace.Name,
Status: workspace.Status,
})
ignoredWorkspaces[workspace.ID] = struct{}{}
}
}
}
Expand Down
84 changes: 62 additions & 22 deletions vpn/tunnel_internal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -293,6 +293,7 @@ func TestUpdater_createPeerUpdate(t *testing.T) {
ctx: ctx,
netLoopDone: make(chan struct{}),
agents: map[uuid.UUID]tailnet.Agent{},
workspaces: map[uuid.UUID]tailnet.Workspace{},
conn: newFakeConn(tailnet.WorkspaceUpdate{}, hsTime),
}

Expand Down Expand Up @@ -727,6 +728,7 @@ func TestProcessFreshState(t *testing.T) {
wsID1 := uuid.New()
wsID2 := uuid.New()
wsID3 := uuid.New()
wsID4 := uuid.New()

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

ws1 := tailnet.Workspace{ID: wsID1, Name: "ws1"}
ws2 := tailnet.Workspace{ID: wsID2, Name: "ws2"}
ws3 := tailnet.Workspace{ID: wsID3, Name: "ws3"}
ws1 := tailnet.Workspace{ID: wsID1, Name: "ws1", Status: proto.Workspace_RUNNING}
ws2 := tailnet.Workspace{ID: wsID2, Name: "ws2", Status: proto.Workspace_RUNNING}
ws3 := tailnet.Workspace{ID: wsID3, Name: "ws3", Status: proto.Workspace_RUNNING}
ws4 := tailnet.Workspace{ID: wsID4, Name: "ws4", Status: proto.Workspace_RUNNING}

initialAgents := map[uuid.UUID]tailnet.Agent{
agentID1: agent1,
agentID2: agent2,
agentID4: agent4,
}
initialWorkspaces := map[uuid.UUID]tailnet.Workspace{
wsID1: ws1,
wsID2: ws2,
}

tests := []struct {
name string
initialAgents map[uuid.UUID]tailnet.Agent
update *tailnet.WorkspaceUpdate
expectedDelete *tailnet.WorkspaceUpdate // We only care about deletions added by the function
name string
initialAgents map[uuid.UUID]tailnet.Agent
initialWorkspaces map[uuid.UUID]tailnet.Workspace
update *tailnet.WorkspaceUpdate
expectedDelete *tailnet.WorkspaceUpdate // We only care about deletions added by the function
}{
{
name: "NoChange",
initialAgents: initialAgents,
name: "NoChange",
initialAgents: initialAgents,
initialWorkspaces: initialWorkspaces,
update: &tailnet.WorkspaceUpdate{
Kind: tailnet.Snapshot,
UpsertedWorkspaces: []*tailnet.Workspace{&ws1, &ws2},
Expand All @@ -770,8 +779,9 @@ func TestProcessFreshState(t *testing.T) {
},
},
{
name: "AgentAdded", // Agent 3 added in update
initialAgents: initialAgents,
name: "AgentAdded", // Agent 3 added in update
initialAgents: initialAgents,
initialWorkspaces: initialWorkspaces,
update: &tailnet.WorkspaceUpdate{
Kind: tailnet.Snapshot,
UpsertedWorkspaces: []*tailnet.Workspace{&ws1, &ws2, &ws3},
Expand All @@ -785,8 +795,9 @@ func TestProcessFreshState(t *testing.T) {
},
},
{
name: "AgentRemovedWorkspaceAlsoRemoved", // Agent 2 removed, ws2 also removed
initialAgents: initialAgents,
name: "AgentRemovedWorkspaceAlsoRemoved", // Agent 2 removed, ws2 also removed
initialAgents: initialAgents,
initialWorkspaces: initialWorkspaces,
update: &tailnet.WorkspaceUpdate{
Kind: tailnet.Snapshot,
UpsertedWorkspaces: []*tailnet.Workspace{&ws1}, // ws2 not present
Expand All @@ -795,15 +806,18 @@ func TestProcessFreshState(t *testing.T) {
DeletedAgents: []*tailnet.Agent{},
},
expectedDelete: &tailnet.WorkspaceUpdate{
DeletedWorkspaces: []*tailnet.Workspace{{ID: wsID2}}, // Expect ws2 to be deleted
DeletedWorkspaces: []*tailnet.Workspace{
{ID: wsID2, Name: "ws2", Status: proto.Workspace_RUNNING},
}, // Expect ws2 to be deleted
DeletedAgents: []*tailnet.Agent{ // Expect agent2 to be deleted
{ID: agentID2, Name: "agent2", WorkspaceID: wsID2},
},
},
},
{
name: "AgentRemovedWorkspaceStays", // Agent 4 removed, but ws1 stays (due to agent1)
initialAgents: initialAgents,
name: "AgentRemovedWorkspaceStays", // Agent 4 removed, but ws1 stays (due to agent1)
initialAgents: initialAgents,
initialWorkspaces: initialWorkspaces,
update: &tailnet.WorkspaceUpdate{
Kind: tailnet.Snapshot,
UpsertedWorkspaces: []*tailnet.Workspace{&ws1, &ws2}, // ws1 still present
Expand All @@ -819,8 +833,9 @@ func TestProcessFreshState(t *testing.T) {
},
},
{
name: "InitialAgentsEmpty",
initialAgents: map[uuid.UUID]tailnet.Agent{}, // Start with no agents known
name: "InitialAgentsEmpty",
initialAgents: map[uuid.UUID]tailnet.Agent{}, // Start with no agents known
initialWorkspaces: map[uuid.UUID]tailnet.Workspace{},
update: &tailnet.WorkspaceUpdate{
Kind: tailnet.Snapshot,
UpsertedWorkspaces: []*tailnet.Workspace{&ws1, &ws2},
Expand All @@ -834,8 +849,9 @@ func TestProcessFreshState(t *testing.T) {
},
},
{
name: "UpdateEmpty", // Fresh state says nothing exists
initialAgents: initialAgents,
name: "UpdateEmpty", // Snapshot says nothing exists
initialAgents: initialAgents,
initialWorkspaces: initialWorkspaces,
update: &tailnet.WorkspaceUpdate{
Kind: tailnet.Snapshot,
UpsertedWorkspaces: []*tailnet.Workspace{},
Expand All @@ -844,14 +860,35 @@ func TestProcessFreshState(t *testing.T) {
DeletedAgents: []*tailnet.Agent{},
},
expectedDelete: &tailnet.WorkspaceUpdate{ // Expect all initial agents/workspaces to be deleted
DeletedWorkspaces: []*tailnet.Workspace{{ID: wsID1}, {ID: wsID2}}, // ws1 and ws2 deleted
DeletedWorkspaces: []*tailnet.Workspace{
{ID: wsID1, Name: "ws1", Status: proto.Workspace_RUNNING},
{ID: wsID2, Name: "ws2", Status: proto.Workspace_RUNNING},
}, // ws1 and ws2 deleted
DeletedAgents: []*tailnet.Agent{ // agent1, agent2, agent4 deleted
{ID: agentID1, Name: "agent1", WorkspaceID: wsID1},
{ID: agentID2, Name: "agent2", WorkspaceID: wsID2},
{ID: agentID4, Name: "agent4", WorkspaceID: wsID1},
},
},
},
{
name: "WorkspaceWithNoAgents", // Snapshot says nothing exists
initialAgents: initialAgents,
initialWorkspaces: map[uuid.UUID]tailnet.Workspace{wsID1: ws1, wsID2: ws2, wsID4: ws4}, // ws4 has no agents
update: &tailnet.WorkspaceUpdate{
Kind: tailnet.Snapshot,
UpsertedWorkspaces: []*tailnet.Workspace{&ws1, &ws2},
UpsertedAgents: []*tailnet.Agent{&agent1, &agent2, &agent4},
DeletedWorkspaces: []*tailnet.Workspace{},
DeletedAgents: []*tailnet.Agent{},
},
expectedDelete: &tailnet.WorkspaceUpdate{ // Expect all initial agents/workspaces to be deleted
DeletedWorkspaces: []*tailnet.Workspace{
{ID: wsID4, Name: "ws4", Status: proto.Workspace_RUNNING},
}, // ws4 should be deleted
DeletedAgents: []*tailnet.Agent{},
},
},
}

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

processSnapshotUpdate(tt.update, agentsCopy)
workspaceCopy := make(map[uuid.UUID]tailnet.Workspace)
maps.Copy(workspaceCopy, tt.initialWorkspaces)

processSnapshotUpdate(tt.update, agentsCopy, workspaceCopy)

require.ElementsMatch(t, tt.expectedDelete.DeletedAgents, tt.update.DeletedAgents, "DeletedAgents mismatch")
require.ElementsMatch(t, tt.expectedDelete.DeletedWorkspaces, tt.update.DeletedWorkspaces, "DeletedWorkspaces mismatch")
Expand Down
Loading