Skip to content

Commit 1fcd8d7

Browse files
ethanndicksonethanndickson
committed
fix race
1 parent 5231bef commit 1fcd8d7

File tree

5 files changed

+70
-11
lines changed

5 files changed

+70
-11
lines changed

‎coderd/database/dbmem/dbmem.go

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12249,8 +12249,15 @@ func (q *FakeQuerier) UpsertConnectionLog(_ context.Context, arg database.Upsert
1224912249
return q.connectionLogs[i], nil
1225012250
}
1225112251
// Update existing connection with close time and reason
12252-
q.connectionLogs[i].CloseTime = sql.NullTime{Valid: true, Time: arg.Time}
12253-
q.connectionLogs[i].CloseReason = arg.CloseReason
12252+
if !q.connectionLogs[i].CloseTime.Valid {
12253+
q.connectionLogs[i].CloseTime = sql.NullTime{Valid: true, Time: arg.Time}
12254+
}
12255+
if !q.connectionLogs[i].CloseReason.Valid {
12256+
q.connectionLogs[i].CloseReason = arg.CloseReason
12257+
}
12258+
if !q.connectionLogs[i].Code.Valid {
12259+
q.connectionLogs[i].Code = arg.Code
12260+
}
1225412261
return q.connectionLogs[i], nil
1225512262
}
1225612263
}

‎coderd/database/querier_test.go

Lines changed: 48 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2362,7 +2362,7 @@ func TestUpsertConnectionLog(t *testing.T) {
23622362
require.Equal(t, log, rows[0].ConnectionLog)
23632363
})
23642364

2365-
t.Run("NoConnect", func(t *testing.T) {
2365+
t.Run("DisconnectThenConnect", func(t *testing.T) {
23662366
t.Parallel()
23672367

23682368
db, _ := dbtestutil.NewDB(t)
@@ -2375,7 +2375,7 @@ func TestUpsertConnectionLog(t *testing.T) {
23752375

23762376
// Insert just a 'disconect' event
23772377
disconnectTime := dbtime.Now()
2378-
connectParams := database.UpsertConnectionLogParams{
2378+
disconnectParams := database.UpsertConnectionLogParams{
23792379
ID: uuid.New(),
23802380
Time: disconnectTime,
23812381
OrganizationID: ws.OrganizationID,
@@ -2386,20 +2386,60 @@ func TestUpsertConnectionLog(t *testing.T) {
23862386
Type: database.ConnectionTypeSsh,
23872387
ConnectionID: uuid.NullUUID{UUID: connectionID, Valid: true},
23882388
ConnectionStatus: database.ConnectionStatusDisconnected,
2389+
CloseReason: sql.NullString{String: "server shutting down", Valid: true},
23892390
}
23902391

2391-
_, err := db.UpsertConnectionLog(ctx, connectParams)
2392+
_, err := db.UpsertConnectionLog(ctx, disconnectParams)
23922393
require.NoError(t, err)
23932394

2394-
rows, err := db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{})
2395+
firstRows, err := db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{})
23952396
require.NoError(t, err)
2396-
require.Len(t, rows, 1)
2397+
require.Len(t, firstRows, 1)
23972398

23982399
// We expect the connection event to be marked as closed with the start
23992400
// and close time being the same.
2400-
require.True(t, rows[0].ConnectionLog.CloseTime.Valid)
2401-
require.Equal(t, disconnectTime, rows[0].ConnectionLog.CloseTime.Time.UTC())
2402-
require.Equal(t, rows[0].ConnectionLog.Time.UTC(), rows[0].ConnectionLog.CloseTime.Time.UTC())
2401+
require.True(t, firstRows[0].ConnectionLog.CloseTime.Valid)
2402+
require.Equal(t, disconnectTime, firstRows[0].ConnectionLog.CloseTime.Time.UTC())
2403+
require.Equal(t, firstRows[0].ConnectionLog.Time.UTC(), firstRows[0].ConnectionLog.CloseTime.Time.UTC())
2404+
2405+
// Now insert a 'connect' event for the same connection.
2406+
// This should be a no op
2407+
connectTime := disconnectTime.Add(time.Second)
2408+
connectParams := database.UpsertConnectionLogParams{
2409+
ID: uuid.New(),
2410+
Time: connectTime,
2411+
OrganizationID: ws.OrganizationID,
2412+
WorkspaceOwnerID: ws.OwnerID,
2413+
WorkspaceID: ws.ID,
2414+
WorkspaceName: ws.Name,
2415+
AgentName: agentName,
2416+
Type: database.ConnectionTypeSsh,
2417+
ConnectionID: uuid.NullUUID{UUID: connectionID, Valid: true},
2418+
ConnectionStatus: database.ConnectionStatusConnected,
2419+
CloseReason: sql.NullString{String: "reconnected", Valid: true},
2420+
Code: sql.NullInt32{Int32: 0, Valid: false},
2421+
}
2422+
2423+
_, err = db.UpsertConnectionLog(ctx, connectParams)
2424+
require.NoError(t, err)
2425+
2426+
secondRows, err := db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{})
2427+
require.NoError(t, err)
2428+
require.Len(t, secondRows, 1)
2429+
require.Equal(t, firstRows, secondRows)
2430+
2431+
// Upsert a disconnection, which should also be a no op
2432+
disconnectParams.CloseReason = sql.NullString{
2433+
String: "updated close reason",
2434+
Valid: true,
2435+
}
2436+
_, err = db.UpsertConnectionLog(ctx, disconnectParams)
2437+
require.NoError(t, err)
2438+
thirdRows, err := db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{})
2439+
require.NoError(t, err)
2440+
require.Len(t, secondRows, 1)
2441+
// The close reason shouldn't be updated
2442+
require.Equal(t, secondRows, thirdRows)
24032443
})
24042444
}
24052445

‎coderd/database/queries.sql.go

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

‎coderd/database/queries/connectionlogs.sql

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,16 +78,22 @@ DO UPDATE SET
7878
-- No-op if the connection is still open.
7979
close_time = CASE
8080
WHEN @connection_status::connection_status = 'disconnected'
81+
-- Can only be set once
82+
AND connection_logs.close_time IS NULL
8183
THEN EXCLUDED."time"
8284
ELSE connection_logs.close_time
8385
END,
8486
close_reason = CASE
8587
WHEN @connection_status::connection_status = 'disconnected'
88+
-- Can only be set once
89+
AND connection_logs.close_reason IS NULL
8690
THEN EXCLUDED.close_reason
8791
ELSE connection_logs.close_reason
8892
END,
8993
code = CASE
9094
WHEN @connection_status::connection_status = 'disconnected'
95+
-- Can only be set once
96+
AND connection_logs.code IS NULL
9197
THEN EXCLUDED.code
9298
ELSE connection_logs.code
9399
END

‎coderd/workspaceapps/db_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -320,7 +320,7 @@ func Test_ResolveRequest(t *testing.T) {
320320
require.Len(t, connLogger.ConnectionLogs(), 1)
321321

322322
var parsedToken workspaceapps.SignedToken
323-
err = jwtutils.Verify(ctx, api.AppSigningKeyCache, cookie.Value, &parsedToken)
323+
err := jwtutils.Verify(ctx, api.AppSigningKeyCache, cookie.Value, &parsedToken)
324324
require.NoError(t, err)
325325
// normalize expiry
326326
require.WithinDuration(t, token.Expiry.Time(), parsedToken.Expiry.Time(), 2*time.Second)

0 commit comments

Comments
 (0)