Function NewMultiReplicaSubscribeFn

enterprise/coderd/x/chatd/chatd.go:149–628

NewMultiReplicaSubscribeFn returns a SubscribeFn that manages relay connections to remote replicas and returns relay message_part events only. OSS handles pubsub subscription, message catch-up, queue updates, status forwarding, and local parts merging. nolint:gocognit // Complexity is inherent to t

(
	cfg MultiReplicaSubscribeConfig,
)

//
//nolint:gocognit // Complexity is inherent to the multi-source merge loop.
func NewMultiReplicaSubscribeFn(
	cfg MultiReplicaSubscribeConfig,
) osschatd.SubscribeFn {
	return func(ctx context.Context, params osschatd.SubscribeFnParams) <-chan codersdk.ChatStreamEvent {
		chatID := params.ChatID
		requestHeader := params.RequestHeader
		logger := params.Logger

		var relayCancel func()
		var relayParts <-chan codersdk.ChatStreamEvent

		// If the chat is currently running on a different worker
		// and we have a remote parts provider, open an initial
		// relay synchronously so the caller gets in-flight
		// message_part events right away.
		var initialRelaySnapshot []codersdk.ChatStreamEvent
		if params.Chat.Status == database.ChatStatusRunning &&
			params.Chat.WorkerID.Valid &&
			params.Chat.WorkerID.UUID != params.WorkerID &&
			cfg.dial() != nil {
			snapshot, parts, cancel, err := cfg.dial()(ctx, chatID, params.Chat.WorkerID.UUID, requestHeader)
			if err == nil {
				relayCancel = cancel
				relayParts = parts
				// Collect relay message_parts to forward at the
				// start of the merge goroutine.
				for _, event := range snapshot {
					if event.Type == codersdk.ChatStreamEventTypeMessagePart {
						initialRelaySnapshot = append(initialRelaySnapshot, event)
					}
				}
			} else {
				logger.Warn(ctx, "failed to open initial relay for chat stream",
					slog.F("chat_id", chatID),
					slog.Error(err),
				)
			}
		}

		// Merge all event sources.
		mergedEvents := make(chan codersdk.ChatStreamEvent, 128)
		// Channel for async relay establishment.
		type relayResult struct {
			parts    <-chan codersdk.ChatStreamEvent
			cancel   func()
			workerID uuid.UUID // the worker this dial targeted
			// err and parts are mutually exclusive: success sets
			// parts; failure sets err (unwrap to *RelayDialError
			// for classification).
			err error
		}
		relayReadyCh := make(chan relayResult, 4)

		// Reset on successful dial or when the relay target
		// changes, so a fresh target starts at the floor delay.
		retryState := newRelayRetryState()
		// Per-dial context so in-flight dials can be canceled when
		// a new dial is initiated or the relay is closed.