55 "fmt"
66 "net/http"
77 "strings"
8+ "sync"
9+ "time"
810
911 libshare "github.com/celestiaorg/go-square/v3/share"
1012 "github.com/filecoin-project/go-jsonrpc"
@@ -13,14 +15,28 @@ import (
1315
1416// Client dials the celestia-node RPC "blob" and "header" namespaces.
1517type Client struct {
16- Blob BlobAPI
17- Header HeaderAPI
18- closer jsonrpc.ClientCloser
18+ Blob BlobAPI
19+ Header HeaderAPI
20+ IsWebSocket bool
21+
22+ mu sync.Mutex
23+ closer jsonrpc.ClientCloser
24+ retryCancel context.CancelFunc // stops the background WS retry loop
1925}
2026
21- // Close closes the underlying JSON-RPC connection.
27+ // Close closes the underlying JSON-RPC connection and stops any
28+ // background WebSocket retry loop.
2229func (c * Client ) Close () {
23- if c != nil && c .closer != nil {
30+ if c == nil {
31+ return
32+ }
33+ c .mu .Lock ()
34+ if c .retryCancel != nil {
35+ c .retryCancel ()
36+ c .retryCancel = nil
37+ }
38+ c .mu .Unlock ()
39+ if c .closer != nil {
2440 c .closer ()
2541 }
2642}
@@ -72,18 +88,88 @@ func NewClient(ctx context.Context, addr, token string, authHeaderName string) (
7288// NewWSClient connects to the DA RPC endpoint over WebSocket.
7389// Automatically converts http:// to ws:// (and https:// to wss://).
7490// Supports channel-based subscriptions (e.g. Subscribe).
75- // Note: WebSocket connections are eager — they connect at creation time
76- // if the initial WS dial fails, falls back to HTTP polling for the entire session.
91+ // WebSocket connections are eager — they connect at creation time.
92+ // If the initial WS dial fails, it falls back to HTTP polling and spawns a
93+ // background goroutine that periodically retries the WS connection. When
94+ // the WS endpoint becomes reachable, the transport is transparently upgraded.
7795func NewWSClient (ctx context.Context , logger zerolog.Logger , addr , token string , authHeaderName string ) (* Client , error ) {
7896 client , err := NewClient (ctx , httpToWS (addr ), token , authHeaderName )
7997 if err != nil {
8098 logger .Warn ().Err (err ).Msg ("DA websocket connection failed, falling back to DA polling" )
81- return NewClient (ctx , addr , token , authHeaderName )
99+ client , err = NewClient (ctx , addr , token , authHeaderName )
100+ if err != nil {
101+ return nil , err
102+ }
103+ client .IsWebSocket = false
104+
105+ // Retry WS in the background so transient outages don't force a permanent downgrade.
106+ retryCtx , retryCancel := context .WithCancel (context .Background ())
107+ client .retryCancel = retryCancel
108+ go client .retryWSLoop (retryCtx , logger , addr , token , authHeaderName )
109+
110+ return client , nil
82111 }
83112
113+ client .IsWebSocket = true
84114 return client , nil
85115}
86116
117+ const wsRetryInterval = 30 * time .Second
118+
119+ // retryWSLoop periodically attempts to re-establish a WebSocket connection.
120+ // When successful, it swaps the transport in-place and exits.
121+ func (c * Client ) retryWSLoop (ctx context.Context , logger zerolog.Logger , addr , token , authHeaderName string ) {
122+ ticker := time .NewTicker (wsRetryInterval )
123+ defer ticker .Stop ()
124+
125+ for {
126+ select {
127+ case <- ctx .Done ():
128+ return
129+ case <- ticker .C :
130+ if c .tryUpgradeWS (ctx , logger , addr , token , authHeaderName ) {
131+ return
132+ }
133+ }
134+ }
135+ }
136+
137+ // tryUpgradeWS attempts to open a WS connection and, if successful, swaps
138+ // the transport internals so subsequent calls use WebSocket. Returns true
139+ // when the upgrade succeeds (or the client is already on WS).
140+ func (c * Client ) tryUpgradeWS (ctx context.Context , logger zerolog.Logger , addr , token , authHeaderName string ) bool {
141+ wsClient , err := NewClient (ctx , httpToWS (addr ), token , authHeaderName )
142+ if err != nil {
143+ return false
144+ }
145+
146+ c .mu .Lock ()
147+ defer c .mu .Unlock ()
148+
149+ // Another goroutine may have already upgraded.
150+ if c .IsWebSocket {
151+ wsClient .Close ()
152+ return true
153+ }
154+
155+ // Swap function pointers from the new WS client into the active client.
156+ c .Blob .Internal = wsClient .Blob .Internal
157+ c .Header .Internal = wsClient .Header .Internal
158+
159+ // Close the old HTTP connections and wire the new closer.
160+ oldCloser := c .closer
161+ c .closer = func () {
162+ wsClient .closer ()
163+ if oldCloser != nil {
164+ oldCloser ()
165+ }
166+ }
167+
168+ c .IsWebSocket = true
169+ logger .Info ().Msg ("DA websocket connection restored, switching back from HTTP polling" )
170+ return true
171+ }
172+
87173// BlobAPI mirrors celestia-node's blob module (nodebuilder/blob/blob.go).
88174// jsonrpc.NewClient wires Internal.* to RPC stubs.
89175type BlobAPI struct {
0 commit comments