dialRelay opens a WebSocket to the replica owning chatID and returns any buffered message_part snapshot plus a live channel of subsequent events. Handshake failures return an error unwrapping to *RelayDialError so callers can classify via IsUnrecoverable. websocket.Dial is called directly (not via
( ctx context.Context, chatID uuid.UUID, workerID uuid.UUID, requestHeader http.Header, cfg MultiReplicaSubscribeConfig, clk quartz.Clock, )
| 682 | // websocket.Dial is called directly (not via the SDK wrapper) so we |
| 683 | // can read *http.Response.StatusCode for classification. |
| 684 | func dialRelay( |
| 685 | ctx context.Context, |
| 686 | chatID uuid.UUID, |
| 687 | workerID uuid.UUID, |
| 688 | requestHeader http.Header, |
| 689 | cfg MultiReplicaSubscribeConfig, |
| 690 | clk quartz.Clock, |
| 691 | ) ( |
| 692 | snapshot []codersdk.ChatStreamEvent, |
| 693 | parts <-chan codersdk.ChatStreamEvent, |
| 694 | cancel func(), |
| 695 | err error, |
| 696 | ) { |
| 697 | address, ok := cfg.ResolveReplicaAddress(ctx, workerID) |
| 698 | if !ok { |
| 699 | return nil, nil, nil, &RelayDialError{ |
| 700 | Err: xerrors.New("dial relay stream: worker replica not found"), |
| 701 | } |
| 702 | } |
| 703 | |
| 704 | wsURL, err := buildRelayURL(address, chatID) |
| 705 | if err != nil { |
| 706 | return nil, nil, nil, &RelayDialError{ |
| 707 | Err: xerrors.Errorf("dial relay stream: %w", err), |
| 708 | } |
| 709 | } |
| 710 | |
| 711 | replicaID := cfg.ReplicaIDFn() |
| 712 | headers := make(http.Header, 2) |
| 713 | headers.Set(codersdk.SessionTokenHeader, extractSessionToken(requestHeader)) |
| 714 | headers.Set(RelaySourceHeader, replicaID.String()) |
| 715 | |
| 716 | relayCtx, relayCancel := context.WithCancel(ctx) |
| 717 | conn, resp, dialErr := websocket.Dial(relayCtx, wsURL, &websocket.DialOptions{ |
| 718 | HTTPClient: cfg.ReplicaHTTPClient, |
| 719 | HTTPHeader: headers, |
| 720 | CompressionMode: websocket.CompressionDisabled, |
| 721 | }) |
| 722 | status := 0 |
| 723 | if resp != nil { |
| 724 | status = resp.StatusCode |
| 725 | // The websocket library closes resp.Body on success; on |
| 726 | // failure we close it ourselves so we don't leak the TCP |
| 727 | // connection. |
| 728 | if dialErr != nil && resp.Body != nil { |
| 729 | _ = resp.Body.Close() |
| 730 | } |
| 731 | } |
| 732 | if dialErr != nil { |
| 733 | relayCancel() |
| 734 | return nil, nil, nil, &RelayDialError{ |
| 735 | HTTPStatus: status, |
| 736 | Err: xerrors.Errorf("dial relay stream: %w", dialErr), |
| 737 | } |
| 738 | } |
| 739 | // Match the server's 4 MiB read limit in codersdk.StreamChat so |
| 740 | // large message_part batches don't trip the default 32 KiB cap. |
| 741 | conn.SetReadLimit(1 << 22) |
no test coverage detected