fix race

ethanndickson · ethanndickson · commit 1fcd8d7261cd · 2025-06-30T11:55:38.000Z
diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -12249,8 +12249,15 @@ func (q *FakeQuerier) UpsertConnectionLog(_ context.Context, arg database.Upsert
 				return q.connectionLogs[i], nil
 			}
 			// Update existing connection with close time and reason
-			q.connectionLogs[i].CloseTime = sql.NullTime{Valid: true, Time: arg.Time}
-			q.connectionLogs[i].CloseReason = arg.CloseReason
+			if !q.connectionLogs[i].CloseTime.Valid {
+				q.connectionLogs[i].CloseTime = sql.NullTime{Valid: true, Time: arg.Time}
+			}
+			if !q.connectionLogs[i].CloseReason.Valid {
+				q.connectionLogs[i].CloseReason = arg.CloseReason
+			}
+			if !q.connectionLogs[i].Code.Valid {
+				q.connectionLogs[i].Code = arg.Code
+			}
 			return q.connectionLogs[i], nil
 		}
 	}
diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go
@@ -2362,7 +2362,7 @@ func TestUpsertConnectionLog(t *testing.T) {
 		require.Equal(t, log, rows[0].ConnectionLog)
 	})
 
-	t.Run("NoConnect", func(t *testing.T) {
+	t.Run("DisconnectThenConnect", func(t *testing.T) {
 		t.Parallel()
 
 		db, _ := dbtestutil.NewDB(t)
@@ -2375,7 +2375,7 @@ func TestUpsertConnectionLog(t *testing.T) {
 
 		// Insert just a 'disconect' event
 		disconnectTime := dbtime.Now()
-		connectParams := database.UpsertConnectionLogParams{
+		disconnectParams := database.UpsertConnectionLogParams{
 			ID:               uuid.New(),
 			Time:             disconnectTime,
 			OrganizationID:   ws.OrganizationID,
@@ -2386,20 +2386,60 @@ func TestUpsertConnectionLog(t *testing.T) {
 			Type:             database.ConnectionTypeSsh,
 			ConnectionID:     uuid.NullUUID{UUID: connectionID, Valid: true},
 			ConnectionStatus: database.ConnectionStatusDisconnected,
+			CloseReason:      sql.NullString{String: "server shutting down", Valid: true},
 		}
 
-		_, err := db.UpsertConnectionLog(ctx, connectParams)
+		_, err := db.UpsertConnectionLog(ctx, disconnectParams)
 		require.NoError(t, err)
 
-		rows, err := db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{})
+		firstRows, err := db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{})
 		require.NoError(t, err)
-		require.Len(t, rows, 1)
+		require.Len(t, firstRows, 1)
 
 		// We expect the connection event to be marked as closed with the start
 		// and close time being the same.
-		require.True(t, rows[0].ConnectionLog.CloseTime.Valid)
-		require.Equal(t, disconnectTime, rows[0].ConnectionLog.CloseTime.Time.UTC())
-		require.Equal(t, rows[0].ConnectionLog.Time.UTC(), rows[0].ConnectionLog.CloseTime.Time.UTC())
+		require.True(t, firstRows[0].ConnectionLog.CloseTime.Valid)
+		require.Equal(t, disconnectTime, firstRows[0].ConnectionLog.CloseTime.Time.UTC())
+		require.Equal(t, firstRows[0].ConnectionLog.Time.UTC(), firstRows[0].ConnectionLog.CloseTime.Time.UTC())
+
+		// Now insert a 'connect' event for the same connection.
+		// This should be a no op
+		connectTime := disconnectTime.Add(time.Second)
+		connectParams := database.UpsertConnectionLogParams{
+			ID:               uuid.New(),
+			Time:             connectTime,
+			OrganizationID:   ws.OrganizationID,
+			WorkspaceOwnerID: ws.OwnerID,
+			WorkspaceID:      ws.ID,
+			WorkspaceName:    ws.Name,
+			AgentName:        agentName,
+			Type:             database.ConnectionTypeSsh,
+			ConnectionID:     uuid.NullUUID{UUID: connectionID, Valid: true},
+			ConnectionStatus: database.ConnectionStatusConnected,
+			CloseReason:      sql.NullString{String: "reconnected", Valid: true},
+			Code:             sql.NullInt32{Int32: 0, Valid: false},
+		}
+
+		_, err = db.UpsertConnectionLog(ctx, connectParams)
+		require.NoError(t, err)
+
+		secondRows, err := db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{})
+		require.NoError(t, err)
+		require.Len(t, secondRows, 1)
+		require.Equal(t, firstRows, secondRows)
+
+		// Upsert a disconnection, which should also be a no op
+		disconnectParams.CloseReason = sql.NullString{
+			String: "updated close reason",
+			Valid:  true,
+		}
+		_, err = db.UpsertConnectionLog(ctx, disconnectParams)
+		require.NoError(t, err)
+		thirdRows, err := db.GetConnectionLogsOffset(ctx, database.GetConnectionLogsOffsetParams{})
+		require.NoError(t, err)
+		require.Len(t, secondRows, 1)
+		// The close reason shouldn't be updated
+		require.Equal(t, secondRows, thirdRows)
 	})
 }
 
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/connectionlogs.sql b/coderd/database/queries/connectionlogs.sql
@@ -78,16 +78,22 @@ DO UPDATE SET
 	-- No-op if the connection is still open.
 	close_time = CASE
 		WHEN @connection_status::connection_status = 'disconnected'
+		-- Can only be set once
+		AND connection_logs.close_time IS NULL
 		THEN EXCLUDED."time"
 		ELSE connection_logs.close_time
 	END,
 	close_reason = CASE
 		WHEN @connection_status::connection_status = 'disconnected'
+		-- Can only be set once
+		AND connection_logs.close_reason IS NULL
 		THEN EXCLUDED.close_reason
 		ELSE connection_logs.close_reason
 	END,
 	code = CASE
 		WHEN @connection_status::connection_status = 'disconnected'
+		-- Can only be set once
+		AND connection_logs.code IS NULL
 		THEN EXCLUDED.code
 		ELSE connection_logs.code
 	END
diff --git a/coderd/workspaceapps/db_test.go b/coderd/workspaceapps/db_test.go
@@ -320,7 +320,7 @@ func Test_ResolveRequest(t *testing.T) {
 					require.Len(t, connLogger.ConnectionLogs(), 1)
 
 					var parsedToken workspaceapps.SignedToken
-					err = jwtutils.Verify(ctx, api.AppSigningKeyCache, cookie.Value, &parsedToken)
+					err := jwtutils.Verify(ctx, api.AppSigningKeyCache, cookie.Value, &parsedToken)
 					require.NoError(t, err)
 					// normalize expiry
 					require.WithinDuration(t, token.Expiry.Time(), parsedToken.Expiry.Time(), 2*time.Second)

Original file line number	Diff line number	Diff line change
`@@ -12249,8 +12249,15 @@ func (q *FakeQuerier) UpsertConnectionLog(_ context.Context, arg database.Upsert`
`12249`	`12249`	`return q.connectionLogs[i], nil`
`12250`	`12250`	`}`
`12251`	`12251`	`// 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`	`+ }`
`12254`	`12261`	`return q.connectionLogs[i], nil`
`12255`	`12262`	`}`
`12256`	`12263`	`}`