mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
fix(chatd): fix relay race conditions, extract enterprise relay logic, move pubsub to OSS (#22589)
## Summary Fixes a bug where interrupting a streaming chat and sending a new message left the relay connected to the wrong replica. Expanded into a broader refactor that cleanly separates concerns: - **OSS** owns pubsub subscription, message catch-up, queue updates, status forwarding, and local parts merging. - **Enterprise** (`enterprise/coderd/chatd`) only manages relay dialing, reconnection, and stale-dial discarding for cross-replica streaming. ## Architecture ### OSS `coderd/chatd/chatd.go` `Subscribe()` builds the initial snapshot then runs a single merge goroutine that handles: - Pubsub subscription for durable events (status, messages, queue, errors) - Message catch-up via `AfterMessageID` - Local `message_part` forwarding - Relay events from enterprise (when `SubscribeFn` is set) - Sends `StatusNotification` to enterprise so it can manage relay lifecycle Key types: - `SubscribeFn` — enterprise hook, returns relay-only events channel - `SubscribeFnParams` — `ChatID`, `Chat`, `WorkerID`, `StatusNotifications`, `RequestHeader`, `DB`, `Logger` - `StatusNotification` — `Status` + `WorkerID`, sent to enterprise on pubsub status changes ### Enterprise `enterprise/coderd/chatd/chatd.go` `NewMultiReplicaSubscribeFn(cfg MultiReplicaSubscribeConfig)` returns a `SubscribeFn` that: - Opens an initial synchronous relay if the chat is running on a remote worker - Reads `StatusNotifications` from OSS to open/close relay connections - Handles async dial, reconnect timers, stale-dial discarding - Returns only relay `message_part` events ## Bug fixes ### Original bug: stale relay dial after interrupt `openRelayAsync` goroutines used `mergedCtx` (subscription-level), not a per-dial context. `closeRelay()` could not cancel in-flight dials. When the user interrupts and a new replica picks up the chat, the old dial goroutine could complete after the new one and deliver a stale `relayResult`. **Fix**: per-dial `dialCtx`/`dialCancel`, `expectedWorkerID` tracking, `workerID` on `relayResult`. `closeRelay()` cancels the dial context and drains `relayReadyCh`. Merge loop rejects mismatched worker IDs. ### Additional fixes - `statusNotifications` send-on-closed-channel race — goroutine now owns `close()` via defer - Enterprise spin-loop on `StatusNotifications` close — two-value receive with nil-out - `hasPubsub` set from `p.pubsub != nil` instead of subscription success — now tracks actual subscription result - `lastMessageID` not initialized from `afterMessageID` — caused duplicate messages on catch-up - `wrappedParts` goroutine leaked remote connection on `dialCtx` cancel - `closeRelay()` did not drain `relayReadyCh` - `setChatWaiting` race with `SendMessage(Interrupt)` — wrapped in `InTx` - `processChat` post-TX side effects fired when chat was taken by another worker — added `errChatTakenByOtherWorker` sentinel - Cancel closure data race on `reconnectTimer` - Bare blocking send on pubsub error path - `localParts` hot-spin after channel close - No-pubsub branch dropped relay events and initial snapshot - Failed relay dial caused permanent stall (no reconnect retry) - DB error during reconnect timer caused permanent stall - `time.NewTimer` replaced with `quartz.Clock` for testable timing ## Tests 9 enterprise tests covering: - Relay reconnect on drop (mock clock) - Async dial does not block merge loop - Relay snapshot delivery - Stale dial discarded after interrupt - Cancel during in-flight dial - Running-to-running worker switch - Failed dial retries (mock clock) - Local worker closes relay - Multiple consecutive reconnects (mock clock) All pass with `-race`.
This commit is contained in:
Kyle Carberry
@@ -6,7 +6,6 @@ import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
@@ -28,7 +27,6 @@ import (
|
||||
"github.com/coder/coder/v2/coderd/database/dbgen"
|
||||
"github.com/coder/coder/v2/coderd/database/dbtestutil"
|
||||
dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub"
|
||||
coderdpubsub "github.com/coder/coder/v2/coderd/pubsub"
|
||||
"github.com/coder/coder/v2/coderd/util/slice"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
"github.com/coder/coder/v2/provisioner/echo"
|
||||
@@ -1133,30 +1131,6 @@ func newTestServer(
|
||||
return server
|
||||
}
|
||||
|
||||
func newTestServerWithRelay(
|
||||
t *testing.T,
|
||||
db database.Store,
|
||||
ps dbpubsub.Pubsub,
|
||||
replicaID uuid.UUID,
|
||||
provider chatd.RemotePartsProvider,
|
||||
) *chatd.Server {
|
||||
t.Helper()
|
||||
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
server := chatd.New(chatd.Config{
|
||||
Logger: logger,
|
||||
Database: db,
|
||||
ReplicaID: replicaID,
|
||||
Pubsub: ps,
|
||||
RemotePartsProvider: provider,
|
||||
PendingChatAcquireInterval: testutil.WaitSuperLong,
|
||||
})
|
||||
t.Cleanup(func() {
|
||||
require.NoError(t, server.Close())
|
||||
})
|
||||
return server
|
||||
}
|
||||
|
||||
func seedChatDependencies(
|
||||
ctx context.Context,
|
||||
t *testing.T,
|
||||
@@ -1213,293 +1187,6 @@ func setOpenAIProviderBaseURL(
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestSubscribeRelayReconnectsOnDrop(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
workerID := uuid.New()
|
||||
subscriberID := uuid.New()
|
||||
|
||||
var callCount atomic.Int32
|
||||
|
||||
provider := func(ctx context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) (
|
||||
[]codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error,
|
||||
) {
|
||||
call := callCount.Add(1)
|
||||
ch := make(chan codersdk.ChatStreamEvent, 10)
|
||||
if call == 1 {
|
||||
// First relay: send a part then close to simulate a drop.
|
||||
ch <- codersdk.ChatStreamEvent{
|
||||
Type: codersdk.ChatStreamEventTypeMessagePart,
|
||||
MessagePart: &codersdk.ChatStreamMessagePart{
|
||||
Role: "assistant",
|
||||
Part: codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeText, Text: "first-relay"},
|
||||
},
|
||||
}
|
||||
close(ch)
|
||||
} else {
|
||||
// Second relay: send a different part, keep open.
|
||||
ch <- codersdk.ChatStreamEvent{
|
||||
Type: codersdk.ChatStreamEventTypeMessagePart,
|
||||
MessagePart: &codersdk.ChatStreamMessagePart{
|
||||
Role: "assistant",
|
||||
Part: codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeText, Text: "second-relay"},
|
||||
},
|
||||
}
|
||||
// Don't close — keep alive so the subscriber stays connected.
|
||||
}
|
||||
return nil, ch, func() {}, nil
|
||||
}
|
||||
|
||||
subscriber := newTestServerWithRelay(t, db, ps, subscriberID, provider)
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
user, model := seedChatDependencies(ctx, t, db)
|
||||
|
||||
// Create a chat and mark it as running on a remote worker.
|
||||
chat, err := subscriber.CreateChat(ctx, chatd.CreateOptions{
|
||||
OwnerID: user.ID,
|
||||
Title: "relay-reconnect",
|
||||
ModelConfigID: model.ID,
|
||||
InitialUserContent: []fantasy.Content{fantasy.TextContent{Text: "hello"}},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
chat, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{
|
||||
ID: chat.ID,
|
||||
Status: database.ChatStatusRunning,
|
||||
WorkerID: uuid.NullUUID{UUID: workerID, Valid: true},
|
||||
StartedAt: sql.NullTime{Time: time.Now(), Valid: true},
|
||||
HeartbeatAt: sql.NullTime{Time: time.Now(), Valid: true},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
_, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0)
|
||||
require.True(t, ok)
|
||||
t.Cleanup(cancel)
|
||||
|
||||
// Should get the first relay part.
|
||||
require.Eventually(t, func() bool {
|
||||
select {
|
||||
case event := <-events:
|
||||
if event.Type == codersdk.ChatStreamEventTypeMessagePart &&
|
||||
event.MessagePart != nil &&
|
||||
event.MessagePart.Part.Text == "first-relay" {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}, testutil.WaitMedium, testutil.IntervalFast)
|
||||
|
||||
// After the first relay closes, a reconnection should happen and
|
||||
// deliver the second relay part.
|
||||
require.Eventually(t, func() bool {
|
||||
select {
|
||||
case event := <-events:
|
||||
if event.Type == codersdk.ChatStreamEventTypeMessagePart &&
|
||||
event.MessagePart != nil &&
|
||||
event.MessagePart.Part.Text == "second-relay" {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}, testutil.WaitMedium, testutil.IntervalFast)
|
||||
|
||||
require.GreaterOrEqual(t, int(callCount.Load()), 2)
|
||||
}
|
||||
|
||||
func TestSubscribeRelayAsyncDoesNotBlock(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
workerID := uuid.New()
|
||||
subscriberID := uuid.New()
|
||||
|
||||
dialStarted := make(chan struct{})
|
||||
dialContinue := make(chan struct{})
|
||||
|
||||
provider := func(ctx context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) (
|
||||
[]codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error,
|
||||
) {
|
||||
// Signal that the dial has started, then block until released.
|
||||
select {
|
||||
case <-dialStarted:
|
||||
default:
|
||||
close(dialStarted)
|
||||
}
|
||||
select {
|
||||
case <-dialContinue:
|
||||
case <-ctx.Done():
|
||||
return nil, nil, nil, ctx.Err()
|
||||
}
|
||||
ch := make(chan codersdk.ChatStreamEvent, 10)
|
||||
return nil, ch, func() {}, nil
|
||||
}
|
||||
|
||||
subscriber := newTestServerWithRelay(t, db, ps, subscriberID, provider)
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
user, model := seedChatDependencies(ctx, t, db)
|
||||
|
||||
// Create a chat in pending status.
|
||||
chat, err := subscriber.CreateChat(ctx, chatd.CreateOptions{
|
||||
OwnerID: user.ID,
|
||||
Title: "relay-async-nonblock",
|
||||
ModelConfigID: model.ID,
|
||||
InitialUserContent: []fantasy.Content{fantasy.TextContent{Text: "hello"}},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Subscribe before the chat is marked running so the relay opens
|
||||
// via pubsub notification (openRelayAsync path).
|
||||
_, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0)
|
||||
require.True(t, ok)
|
||||
t.Cleanup(cancel)
|
||||
|
||||
// Now mark the chat as running on a remote worker. This publishes
|
||||
// a status notification which triggers openRelayAsync on the
|
||||
// subscriber.
|
||||
notify := coderdpubsub.ChatStreamNotifyMessage{
|
||||
Status: string(database.ChatStatusRunning),
|
||||
WorkerID: workerID.String(),
|
||||
}
|
||||
payload, err := json.Marshal(notify)
|
||||
require.NoError(t, err)
|
||||
err = ps.Publish(coderdpubsub.ChatStreamNotifyChannel(chat.ID), payload)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Wait for the relay dial to actually start (blocking in the
|
||||
// provider).
|
||||
select {
|
||||
case <-dialStarted:
|
||||
case <-ctx.Done():
|
||||
t.Fatal("timed out waiting for relay dial to start")
|
||||
}
|
||||
|
||||
// While the relay is still dialing (provider is blocked), publish
|
||||
// another status change. If openRelayAsync blocked the select loop
|
||||
// this event would never arrive.
|
||||
statusNotify := coderdpubsub.ChatStreamNotifyMessage{
|
||||
Status: string(database.ChatStatusWaiting),
|
||||
}
|
||||
statusPayload, err := json.Marshal(statusNotify)
|
||||
require.NoError(t, err)
|
||||
err = ps.Publish(coderdpubsub.ChatStreamNotifyChannel(chat.ID), statusPayload)
|
||||
require.NoError(t, err)
|
||||
|
||||
// The waiting status event should arrive promptly despite the
|
||||
// relay still dialing.
|
||||
require.Eventually(t, func() bool {
|
||||
select {
|
||||
case event := <-events:
|
||||
return event.Type == codersdk.ChatStreamEventTypeStatus &&
|
||||
event.Status != nil &&
|
||||
event.Status.Status == codersdk.ChatStatusWaiting
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}, testutil.WaitShort, testutil.IntervalFast)
|
||||
|
||||
// Unblock the relay dial so the test can clean up.
|
||||
close(dialContinue)
|
||||
}
|
||||
|
||||
func TestSubscribeRelaySnapshotDelivered(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
workerID := uuid.New()
|
||||
subscriberID := uuid.New()
|
||||
|
||||
provider := func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) (
|
||||
[]codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error,
|
||||
) {
|
||||
// Return a non-empty snapshot with two parts.
|
||||
snapshot := []codersdk.ChatStreamEvent{
|
||||
{
|
||||
Type: codersdk.ChatStreamEventTypeMessagePart,
|
||||
MessagePart: &codersdk.ChatStreamMessagePart{
|
||||
Role: "assistant",
|
||||
Part: codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeText, Text: "snap-one"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Type: codersdk.ChatStreamEventTypeMessagePart,
|
||||
MessagePart: &codersdk.ChatStreamMessagePart{
|
||||
Role: "assistant",
|
||||
Part: codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeText, Text: "snap-two"},
|
||||
},
|
||||
},
|
||||
}
|
||||
ch := make(chan codersdk.ChatStreamEvent, 10)
|
||||
// Also send a live part after the snapshot.
|
||||
ch <- codersdk.ChatStreamEvent{
|
||||
Type: codersdk.ChatStreamEventTypeMessagePart,
|
||||
MessagePart: &codersdk.ChatStreamMessagePart{
|
||||
Role: "assistant",
|
||||
Part: codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeText, Text: "live-part"},
|
||||
},
|
||||
}
|
||||
return snapshot, ch, func() {}, nil
|
||||
}
|
||||
|
||||
subscriber := newTestServerWithRelay(t, db, ps, subscriberID, provider)
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
user, model := seedChatDependencies(ctx, t, db)
|
||||
|
||||
// Create a chat already running on a remote worker.
|
||||
chat, err := subscriber.CreateChat(ctx, chatd.CreateOptions{
|
||||
OwnerID: user.ID,
|
||||
Title: "relay-snapshot",
|
||||
ModelConfigID: model.ID,
|
||||
InitialUserContent: []fantasy.Content{fantasy.TextContent{Text: "hello"}},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{
|
||||
ID: chat.ID,
|
||||
Status: database.ChatStatusRunning,
|
||||
WorkerID: uuid.NullUUID{UUID: workerID, Valid: true},
|
||||
StartedAt: sql.NullTime{Time: time.Now(), Valid: true},
|
||||
HeartbeatAt: sql.NullTime{Time: time.Now(), Valid: true},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
initialSnapshot, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0)
|
||||
require.True(t, ok)
|
||||
t.Cleanup(cancel)
|
||||
|
||||
// The initial snapshot should contain the two relay snapshot parts.
|
||||
var snapshotTexts []string
|
||||
for _, event := range initialSnapshot {
|
||||
if event.Type == codersdk.ChatStreamEventTypeMessagePart && event.MessagePart != nil {
|
||||
snapshotTexts = append(snapshotTexts, event.MessagePart.Part.Text)
|
||||
}
|
||||
}
|
||||
require.Contains(t, snapshotTexts, "snap-one")
|
||||
require.Contains(t, snapshotTexts, "snap-two")
|
||||
|
||||
// The live part should arrive on the events channel.
|
||||
require.Eventually(t, func() bool {
|
||||
select {
|
||||
case event := <-events:
|
||||
if event.Type == codersdk.ChatStreamEventTypeMessagePart &&
|
||||
event.MessagePart != nil &&
|
||||
event.MessagePart.Part.Text == "live-part" {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}, testutil.WaitMedium, testutil.IntervalFast)
|
||||
}
|
||||
|
||||
func TestCloseDuringShutdownContextCanceledShouldRetryOnNewReplica(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user