mirror of
https://github.com/coder/coder.git
synced 2026-10-16 14:48:56 +08:00
fix(coderd): compare the project under the root lock before skipping a move
patchChat skipped MoveChatToProject when the chat row it read without a lock already had the requested project. A move back to the current project could then return success while an earlier move to another project was still committing, and the earlier move won. The handler now always calls MoveChatToProject, and SetFamilyProject compares the project after it locks the root row. A root that is already in the project returns with no writes and no events, also while the chat runs.
This commit is contained in:
@@ -2,7 +2,9 @@ package coderd_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strings"
|
||||
@@ -18,6 +20,7 @@ import (
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/dbauthz"
|
||||
"github.com/coder/coder/v2/coderd/database/dbgen"
|
||||
"github.com/coder/coder/v2/coderd/database/dbtestutil"
|
||||
"github.com/coder/coder/v2/coderd/rbac"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chatstate"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
@@ -336,6 +339,58 @@ func TestPatchChatProject(t *testing.T) {
|
||||
requireChatProjectID(ctx, t, client, child.ID, nil)
|
||||
})
|
||||
|
||||
t.Run("MoveBackWaitsForAPendingMove", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// A move back to the current project compares with the locked row,
|
||||
// so it cannot succeed while an earlier move commits after it.
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
db, _, sqlDB := dbtestutil.NewDBWithSQLDB(t)
|
||||
client, api := newChatProjectClientWithAPI(t, func(options *coderdtest.Options) {
|
||||
options.Database = db
|
||||
})
|
||||
firstUser := coderdtest.CreateFirstUser(t, client.Client)
|
||||
_ = createChatModel(t, client)
|
||||
projectA := createChatProject(t, client, firstUser.OrganizationID, "Project A")
|
||||
projectB := createChatProject(t, client, firstUser.OrganizationID, "Project B")
|
||||
chat := createIdleChat(ctx, t, client, api, firstUser.OrganizationID)
|
||||
require.NoError(t, client.UpdateChat(ctx, chat.ID, codersdk.UpdateChatRequest{ProjectID: &projectA.ID}))
|
||||
|
||||
pendingMove, err := sqlDB.BeginTx(ctx, nil)
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() {
|
||||
if err := pendingMove.Rollback(); err != nil && !errors.Is(err, sql.ErrTxDone) {
|
||||
t.Errorf("roll back the pending move: %v", err)
|
||||
}
|
||||
})
|
||||
_, err = pendingMove.ExecContext(ctx, `UPDATE chats SET project_id = $1 WHERE id = $2`, projectB.ID, chat.ID)
|
||||
require.NoError(t, err)
|
||||
|
||||
moveBack := make(chan error, 1)
|
||||
go func() {
|
||||
moveBack <- client.UpdateChat(ctx, chat.ID, codersdk.UpdateChatRequest{ProjectID: &projectA.ID})
|
||||
}()
|
||||
testutil.Eventually(ctx, t, func(ctx context.Context) bool {
|
||||
if len(moveBack) > 0 {
|
||||
return true
|
||||
}
|
||||
var waiting int
|
||||
err := sqlDB.QueryRowContext(ctx, `
|
||||
SELECT COUNT(*)
|
||||
FROM pg_stat_activity
|
||||
WHERE datname = current_database()
|
||||
AND pid <> pg_backend_pid()
|
||||
AND wait_event_type = 'Lock'
|
||||
AND query LIKE '%-- name: GetChatByIDForUpdate%'
|
||||
`).Scan(&waiting)
|
||||
return err == nil && waiting == 1
|
||||
}, testutil.IntervalFast, "the move back must answer or wait for the chat row")
|
||||
require.NoError(t, pendingMove.Commit())
|
||||
|
||||
require.NoError(t, testutil.RequireReceive(ctx, t, moveBack))
|
||||
requireChatProjectID(ctx, t, client, chat.ID, &projectA.ID)
|
||||
})
|
||||
|
||||
t.Run("AnotherOwnersProject", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
@@ -505,6 +560,11 @@ func TestPatchChatProject(t *testing.T) {
|
||||
sdkErr := requireSDKError(t, err, http.StatusConflict)
|
||||
require.Equal(t, "Cannot move an active chat. Interrupt or wait for the chat to finish first.", sdkErr.Message)
|
||||
require.Equal(t, before, chatRowsOfFamily(ctx, t, api, chat.ID), "a refused move changes no row")
|
||||
|
||||
// The chat has no project, so removing it from its project
|
||||
// moves nothing and is not refused.
|
||||
require.NoError(t, client.UpdateChat(ctx, chat.ID, codersdk.UpdateChatRequest{ProjectID: &uuid.Nil}))
|
||||
require.Equal(t, before, chatRowsOfFamily(ctx, t, api, chat.ID), "a move to the current project changes no row")
|
||||
})
|
||||
}
|
||||
})
|
||||
|
||||
+1
-1
@@ -2780,7 +2780,7 @@ func (api *API) patchChat(rw http.ResponseWriter, r *http.Request) {
|
||||
|
||||
// The move is the last write, so a refused move keeps the earlier
|
||||
// writes of the request, as a refused archive does.
|
||||
if projectUpdate != nil && *projectUpdate != chat.ProjectID {
|
||||
if projectUpdate != nil {
|
||||
err := api.chatDaemon.MoveChatToProject(ctx, chat, *projectUpdate)
|
||||
if err != nil {
|
||||
if writeChatInvalidState(ctx, rw, err) {
|
||||
|
||||
@@ -508,7 +508,8 @@ func (p *Server) setChatFamilyArchived(ctx context.Context, chat database.Chat,
|
||||
// MoveChatToProject moves a root chat into a project, or out of its
|
||||
// project when projectID is NULL. Its subagent chats follow because they
|
||||
// use the root's project. While the chat or a subagent chat is not idle,
|
||||
// it returns [chatstate.ErrTransitionNotAllowed].
|
||||
// it returns [chatstate.ErrTransitionNotAllowed], unless the chat is
|
||||
// already in that project, which changes nothing.
|
||||
func (p *Server) MoveChatToProject(ctx context.Context, chat database.Chat, projectID uuid.NullUUID) error {
|
||||
if chat.ID == uuid.Nil {
|
||||
return xerrors.New("chat_id is required")
|
||||
|
||||
@@ -173,7 +173,9 @@ type SetFamilyProjectInput struct {
|
||||
// Like [SetFamilyArchived], it locks the root first and then each
|
||||
// member in GetChatFamilyIDsByRootID order, and it publishes only after
|
||||
// the commit. On success it returns one [database.Chat] per family
|
||||
// member, root first.
|
||||
// member, root first. A root already in the project changes nothing and
|
||||
// returns no chats; the check reads the locked row, so a move back to
|
||||
// the current project waits for a move still in flight.
|
||||
func SetFamilyProject(
|
||||
ctx context.Context,
|
||||
store database.Store,
|
||||
@@ -204,6 +206,9 @@ func SetFamilyProject(
|
||||
if root.ParentChatID.Valid {
|
||||
return ErrChatNotRoot
|
||||
}
|
||||
if root.ProjectID == input.ProjectID {
|
||||
return nil
|
||||
}
|
||||
ids, err := tx.GetChatFamilyIDsByRootID(ctx, input.RootID)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("get chat family: %w", err)
|
||||
|
||||
@@ -359,6 +359,41 @@ func TestSetFamilyProjectLocksRootFirst(t *testing.T) {
|
||||
require.False(t, unchanged.ProjectID.Valid, "subagent chats use the root's project")
|
||||
}
|
||||
|
||||
// TestSetFamilyProjectKeepsTheCurrentProject moves a running root chat to
|
||||
// the project it is already in, which changes and publishes nothing.
|
||||
func TestSetFamilyProjectKeepsTheCurrentProject(t *testing.T) {
|
||||
t.Parallel()
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
user, org, model := seedFamilyDeps(t, db)
|
||||
project := dbgen.ChatProject(t, db, database.ChatProject{
|
||||
OrganizationID: org.ID,
|
||||
OwnerID: user.ID,
|
||||
})
|
||||
root := dbgen.Chat(t, db, database.Chat{
|
||||
OrganizationID: org.ID,
|
||||
OwnerID: user.ID,
|
||||
LastModelConfigID: model.ID,
|
||||
Title: "root",
|
||||
Status: database.ChatStatusRunning,
|
||||
ProjectID: uuid.NullUUID{UUID: project.ID, Valid: true},
|
||||
})
|
||||
before, err := db.GetChatByID(ctx, root.ID)
|
||||
require.NoError(t, err)
|
||||
|
||||
publisher := newRecordingPubsub()
|
||||
chats, err := chatstate.SetFamilyProject(ctx, db, publisher, chatstate.SetFamilyProjectInput{
|
||||
RootID: root.ID,
|
||||
ProjectID: root.ProjectID,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, chats)
|
||||
require.Empty(t, publisher.channels)
|
||||
after, err := db.GetChatByID(ctx, root.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, before, after)
|
||||
}
|
||||
|
||||
// waitForLockWaiter waits until exactly one other session waits on a row
|
||||
// lock while it runs the named sqlc query.
|
||||
func waitForLockWaiter(ctx context.Context, t *testing.T, sqlDB *sql.DB, queryName string) {
|
||||
|
||||
Reference in New Issue
Block a user