websocket: Initial refactor for websocket handling to support context propagation

common/common.go

@@ -81,6 +81,7 @@ var (
 	errCannotSetInvalidTimeout = errors.New("cannot set new HTTP client with timeout that is equal or less than 0")
 	errUserAgentInvalid        = errors.New("cannot set invalid user agent")
 	errHTTPClientInvalid       = errors.New("custom http client cannot be nil")
+	errDuplicateContextKey     = errors.New("duplicate context key")
 )
 
 // NilGuard returns an ErrNilPointer with the type of the first nil argument
@@ -695,3 +696,62 @@ func SetIfZero[T comparable](p *T, def T) bool {
 	*p = def
 	return true
 }
+
+var (
+	contextKeys   []any
+	contextKeysMu sync.RWMutex
+)
+
+// RegisterContextKey registers a key to be captured by FreezeCtx
+func RegisterContextKey(key any) {
+	contextKeysMu.Lock()
+	defer contextKeysMu.Unlock()
+	if !slices.Contains(contextKeys, key) {
+		contextKeys = append(contextKeys, key)
+	}
+}
+
+// FrozenContext holds captured context values
+type FrozenContext map[any]any
+
+// FreezeCtx captures values from the context for registered keys
+func FreezeCtx(ctx context.Context) FrozenContext {
+	contextKeysMu.RLock()
+	defer contextKeysMu.RUnlock()
+
+	values := make(FrozenContext, len(contextKeys))
+	for _, key := range contextKeys {
+		if val := ctx.Value(key); val != nil {
+			values[key] = val
+		}
+	}
+	return values
+}
+
+// ThawCtx creates a new context from the frozen context using context.Background() as parent
+func ThawCtx(fc FrozenContext) (context.Context, error) {
+	return MergeCtx(context.Background(), fc)
+}
+
+// MergeCtx adds the frozen values to an existing context
+func MergeCtx(ctx context.Context, fc FrozenContext) (context.Context, error) {
+	for k := range fc {
+		if ctx.Value(k) != nil {
+			return nil, fmt.Errorf("%w: %q", errDuplicateContextKey, k)
+		}
+	}
+	return &mergeCtx{Context: ctx, frozen: fc}, nil
+}
+
+// mergeCtx is a context that merges values from a frozen context and a parent context.
+type mergeCtx struct {
+	context.Context //nolint:containedctx // Using context.WithValue will nest contexts and cause lookup latency
+	frozen          FrozenContext
+}
+
+func (m *mergeCtx) Value(key any) any {
+	if val, ok := m.frozen[key]; ok {
+		return val
+	}
+	return m.Context.Value(key)
+}

common/common_test.go

@@ -1,6 +1,7 @@
 package common
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"net/http"
@@ -691,3 +692,41 @@ func TestSetIfZero(t *testing.T) {
 	assert.True(t, changed, "SetIfZero should change a zero value")
 	assert.Equal(t, "world", s, "SetIfZero should change a zero value")
 }
+
+func TestContextFunctions(t *testing.T) {
+	t.Parallel()
+
+	type key string
+	const k1 key = "key1"
+	const k2 key = "key2"
+	const k3 key = "key3"
+
+	RegisterContextKey(k1)
+	RegisterContextKey(k2)
+
+	ctx := context.WithValue(context.Background(), k1, "value1")
+	ctx = context.WithValue(ctx, k2, "value2")
+	ctx = context.WithValue(ctx, k3, "value3") // Not registered
+
+	frozen := FreezeCtx(ctx)
+
+	assert.Equal(t, "value1", frozen[k1], "should have captured k1")
+	assert.Equal(t, "value2", frozen[k2], "should have captured k2")
+	assert.Zero(t, frozen[k3], "k3 should not be captured")
+
+	thawed, err := ThawCtx(frozen)
+	require.NoError(t, err)
+	assert.Equal(t, "value1", thawed.Value(k1), "should have k1 after thaw")
+	assert.Equal(t, "value2", thawed.Value(k2), "should have k2 after thaw")
+	assert.Nil(t, thawed.Value(k3), "Thawed context should not have k3")
+
+	ctx2 := context.WithValue(context.Background(), k3, "value3_new")
+	merged, err := MergeCtx(ctx2, frozen)
+	require.NoError(t, err)
+	assert.Equal(t, "value1", merged.Value(k1), "should have k1 from frozen")
+	assert.Equal(t, "value2", merged.Value(k2), "should have k2 from frozen")
+	assert.Equal(t, "value3_new", merged.Value(k3), "should have k3 from parent")
+
+	_, err = MergeCtx(merged, frozen)
+	require.ErrorIs(t, err, errDuplicateContextKey, "must error on duplicate keys")
+}

docs/ADD_NEW_EXCHANGE.md

@@ -693,19 +693,17 @@ func (e *Exchange) WsConnect() error {
 // KeepAuthKeyAlive will continuously send messages to
 // keep the WS auth key active
 func (e *Exchange) KeepAuthKeyAlive(ctx context.Context) {
-    e.Websocket.Wg.Add(1)
     defer e.Websocket.Wg.Done()
-    ticks := time.NewTicker(time.Minute * 30)
     for {
         select {
         case <-e.Websocket.ShutdownC:
-            ticks.Stop()
             return
-        case <-ticks.C:
-            err := e.MaintainWsAuthStreamKey(ctx)
-            if err != nil {
-                e.Websocket.DataHandler <- err
-                log.Warnf(log.ExchangeSys, "%s - Unable to renew auth websocket token, may experience shutdown", e.Name)
+        case <-time.After(time.Minute * 30):
+            if err := e.MaintainWsAuthStreamKey(ctx); err != nil {
+                if errSend := e.Websocket.DataHandler.Send(ctx, err); errSend != nil {
+                    log.Errorf(log.WebsocketMgr, "%s %s: %s %s", e.Name, e.Websocket.Conn.GetURL(), errSend, err)
+                }
+                log.Warnf(log.ExchangeSys, "%s %s: Unable to renew auth websocket token, may experience shutdown", e.Name, e.Websocket.Conn.GetURL())
             }
         }
     }
@@ -817,9 +815,7 @@ Run gocryptotrader with the following settings enabled in config
 ```go
 // wsReadData gets and passes on websocket messages for processing
 func (e *Exchange) wsReadData() {
-    e.Websocket.Wg.Add(1)
     defer e.Websocket.Wg.Done()
-
     for {
         select {
         case <-e.Websocket.ShutdownC:
@@ -829,10 +825,10 @@ func (e *Exchange) wsReadData() {
             if resp.Raw == nil {
                 return
             }
-
-            err := e.wsHandleData(resp.Raw)
-            if err != nil {
-                e.Websocket.DataHandler <- err
+            if err := e.wsHandleData(ctx, resp.Raw); err != nil {
+                if errSend := e.Websocket.DataHandler.Send(ctx, err); errSend != nil {
+                    log.Errorf(log.WebsocketMgr, "%s %s: %s %s", e.Name, e.Websocket.Conn.GetURL(), errSend, err)
+                }
             }
         }
     }
@@ -875,15 +871,15 @@ If a suitable struct does not exist in wshandler, wrapper types are the next pre
         if err := json.Unmarshal(respRaw, &resultData);err != nil {
             return err
         }
-        e.Websocket.DataHandler <- &ticker.Price{
+        return e.Websocket.DataHandler.Send(ctx, &ticker.Price{
             ExchangeName: e.Name,
             Bid:          resultData.Ticker.Bid,
             Ask:          resultData.Ticker.Ask,
             Last:         resultData.Ticker.Last,
             LastUpdated:  resultData.Ticker.Time,
             Pair:         p,
             AssetType:    a,
-        }
+        })
     }
 ```
 
@@ -896,7 +892,7 @@ If neither of those provide a suitable struct to store the data in, the data can
             if err != nil {
                 return err
             }
-      e.Websocket.DataHandler <- resultData.FillsData
+            return e.Websocket.DataHandler.Send(ctx, resultData.FillsData)
 ```
 
 - Data Handling can be tested offline similar to the following example:

engine/rpcserver.go

@@ -2947,7 +2947,7 @@ func (s *RPCServer) WebsocketGetInfo(_ context.Context, r *gctrpc.WebsocketGetIn
 }
 
 // WebsocketSetEnabled enables or disables the websocket client
-func (s *RPCServer) WebsocketSetEnabled(_ context.Context, r *gctrpc.WebsocketSetEnabledRequest) (*gctrpc.GenericResponse, error) {
+func (s *RPCServer) WebsocketSetEnabled(ctx context.Context, r *gctrpc.WebsocketSetEnabledRequest) (*gctrpc.GenericResponse, error) {
 	exch, err := s.GetExchangeByName(r.Exchange)
 	if err != nil {
 		return nil, err
@@ -2964,7 +2964,7 @@ func (s *RPCServer) WebsocketSetEnabled(_ context.Context, r *gctrpc.WebsocketSe
 	}
 
 	if r.Enable {
-		err = w.Enable()
+		err = w.Enable(s.rpcContextToLongLivedSession(ctx))
 		if err != nil {
 			return nil, err
 		}
@@ -3013,7 +3013,7 @@ func (s *RPCServer) WebsocketGetSubscriptions(_ context.Context, r *gctrpc.Webso
 }
 
 // WebsocketSetProxy sets client websocket connection proxy
-func (s *RPCServer) WebsocketSetProxy(_ context.Context, r *gctrpc.WebsocketSetProxyRequest) (*gctrpc.GenericResponse, error) {
+func (s *RPCServer) WebsocketSetProxy(ctx context.Context, r *gctrpc.WebsocketSetProxyRequest) (*gctrpc.GenericResponse, error) {
 	exch, err := s.GetExchangeByName(r.Exchange)
 	if err != nil {
 		return nil, err
@@ -3024,7 +3024,7 @@ func (s *RPCServer) WebsocketSetProxy(_ context.Context, r *gctrpc.WebsocketSetP
 		return nil, fmt.Errorf("websocket not supported for exchange %s", r.Exchange)
 	}
 
-	err = w.SetProxyAddress(r.Proxy)
+	err = w.SetProxyAddress(s.rpcContextToLongLivedSession(ctx), r.Proxy)
 	if err != nil {
 		return nil, err
 	}
@@ -5884,3 +5884,13 @@ func (s *RPCServer) GetCurrencyTradeURL(ctx context.Context, r *gctrpc.GetCurren
 		Url: url,
 	}, nil
 }
+
+// rpcContextToLongLivedSession converts a short-lived incoming context to a long-lived outgoing context, this is due
+// to the incoming context being cancelled when the RPC call completes.
+func (s *RPCServer) rpcContextToLongLivedSession(ctx context.Context) context.Context {
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		md = metadata.New(nil) // Fallback to empty metadata
+	}
+	return metadata.NewOutgoingContext(context.Background(), md)
+}

engine/rpcserver_test.go

@@ -3869,3 +3869,51 @@ func TestGetCurrencyTradeURL(t *testing.T) {
 	assert.NoError(t, err)
 	assert.NotEmpty(t, resp.Url)
 }
+
+func TestRPCContextToLongLivedSession(t *testing.T) {
+	t.Parallel()
+	s := &RPCServer{}
+	t.Run("preserve metadata", func(t *testing.T) {
+		t.Parallel()
+		incomingMD := metadata.Pairs(
+			"authorization", "Bearer token",
+			"x-id", "123",
+		)
+		inCtx, cancel := context.WithCancel(metadata.NewIncomingContext(context.Background(), incomingMD))
+		defer cancel()
+
+		newCtx := s.rpcContextToLongLivedSession(inCtx)
+		require.NotNil(t, newCtx)
+
+		outMD, ok := metadata.FromOutgoingContext(newCtx)
+		require.True(t, ok)
+		require.Len(t, outMD, len(incomingMD))
+		assert.ElementsMatch(t, incomingMD.Get("authorization"), outMD.Get("authorization"))
+		assert.ElementsMatch(t, incomingMD.Get("x-id"), outMD.Get("x-id"))
+
+		cancel()
+		assert.Eventually(t, func() bool {
+			return inCtx.Err() != nil
+		}, time.Second, 10*time.Millisecond)
+		assert.Eventually(t, func() bool {
+			return newCtx.Err() == nil
+		}, 500*time.Millisecond, 20*time.Millisecond)
+	})
+
+	t.Run("no metadata", func(t *testing.T) {
+		t.Parallel()
+		inCtx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		newCtx := s.rpcContextToLongLivedSession(inCtx)
+		require.NotNil(t, newCtx)
+
+		outMD, ok := metadata.FromOutgoingContext(newCtx)
+		require.True(t, ok)
+		assert.Empty(t, outMD)
+
+		cancel()
+		assert.Eventually(t, func() bool { return inCtx.Err() != nil }, time.Second, 10*time.Millisecond)
+		assert.Eventually(t, func() bool { return newCtx.Err() == nil }, 500*time.Millisecond, 20*time.Millisecond)
+	})
+}

engine/websocketroutine_manager.go

@@ -1,6 +1,7 @@
 package engine
 
 import (
+	"context"
 	"fmt"
 	"sync"
 	"sync/atomic"
@@ -139,7 +140,7 @@ func (m *WebsocketRoutineManager) websocketRoutine() {
 				log.Errorf(log.WebsocketMgr, "%v", err)
 			}
 
-			if err := ws.Connect(); err != nil {
+			if err := ws.Connect(context.TODO()); err != nil {
 				log.Errorf(log.WebsocketMgr, "%v", err)
 			}
 		})
@@ -167,14 +168,13 @@ func (m *WebsocketRoutineManager) websocketDataReceiver(ws *websocket.Manager) e
 			select {
 			case <-m.shutdown:
 				return
-			case data := <-ws.ToRoutine:
-				if data == nil {
+			case payload := <-ws.DataHandler.C:
+				if payload.Data == nil {
 					log.Errorf(log.WebsocketMgr, "exchange %s nil data sent to websocket", ws.GetName())
 				}
 				m.mu.RLock()
 				for x := range m.dataHandlers {
-					err := m.dataHandlers[x](ws.GetName(), data)
-					if err != nil {
+					if err := m.dataHandlers[x](ws.GetName(), payload.Data); err != nil {
 						log.Errorln(log.WebsocketMgr, err)
 					}
 				}

engine/websocketroutine_manager_test.go

@@ -258,16 +258,18 @@ func TestRegisterWebsocketDataHandlerWithFunctionality(t *testing.T) {
 	}
 
 	mock := websocket.NewManager()
-	mock.ToRoutine = make(chan any)
 	m.state = readyState
 	err = m.websocketDataReceiver(mock)
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	mock.ToRoutine <- nil
-	mock.ToRoutine <- 1336
-	mock.ToRoutine <- "intercepted"
+	err = mock.DataHandler.Send(t.Context(), nil)
+	require.NoError(t, err)
+	err = mock.DataHandler.Send(t.Context(), 1336)
+	require.NoError(t, err)
+	err = mock.DataHandler.Send(t.Context(), "intercepted")
+	require.NoError(t, err)
 
 	if r := <-dataChan; r != "intercepted" {
 		t.Fatal("unexpected value received")