websocket: Initial refactor for websocket handling to support context propagation

common/common.go

@@ -696,3 +696,58 @@ 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 FreezeContext
+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
+
+// FreezeContext captures values from the context for registered keys
+func FreezeContext(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
+}
+
+// ThawContext creates a new context from the frozen context using context.Background() as parent
+func ThawContext(fc FrozenContext) context.Context {
+	return MergeContext(context.Background(), fc)
+}
+
+// MergeContext adds the frozen values to an existing context
+func MergeContext(ctx context.Context, fc FrozenContext) context.Context {
+	return &mergedContext{Context: ctx, frozen: fc}
+}
+
+// mergedContext is a context that has merged values from a frozen context and a parent context.
+// frozen values are stored in FrozenContext instead of nested context.WithValue because of the performance of calling WithValue N+ times on messages being frozen
+type mergedContext struct {
+	context.Context //nolint:containedctx // mergedContext implements context.Context
+	frozen          FrozenContext
+}
+
+func (m *mergedContext) 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,36 @@ 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 := FreezeContext(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 := ThawContext(frozen)
+	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 := MergeContext(ctx2, frozen)
+	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")
+}

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,11 +2964,9 @@ func (s *RPCServer) WebsocketSetEnabled(_ context.Context, r *gctrpc.WebsocketSe
 	}
 
 	if r.Enable {
-		err = w.Enable()
-		if err != nil {
+		if err := w.Enable(context.WithoutCancel(ctx)); err != nil {
 			return nil, err
 		}
-
 		exchCfg.Features.Enabled.Websocket = true
 		return &gctrpc.GenericResponse{Status: MsgStatusSuccess, Data: "websocket enabled"}, nil
 	}
@@ -3013,7 +3011,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,15 +3022,12 @@ 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)
-	if err != nil {
+	if err := w.SetProxyAddress(context.WithoutCancel(ctx), r.Proxy); err != nil {
 		return nil, err
 	}
 	return &gctrpc.GenericResponse{
 		Status: MsgStatusSuccess,
-		Data: fmt.Sprintf("new proxy has been set [%s] for %s websocket connection",
-			r.Exchange,
-			r.Proxy),
+		Data:   fmt.Sprintf("new proxy has been set [%s] for %s websocket connection", r.Exchange, r.Proxy),
 	}, nil
 }
 

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")

exchange/stream/relay.go

@@ -0,0 +1,48 @@
+package stream
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	"github.com/thrasher-corp/gocryptotrader/common"
+)
+
+var errChannelBufferFull = errors.New("channel buffer is full")
+
+// Relay defines a channel relay for messages
+type Relay struct {
+	C    <-chan Payload
+	comm chan Payload
+}
+
+// Payload represents a relayed message with a context
+type Payload struct {
+	Ctx  common.FrozenContext
+	Data any
+}
+
+// NewRelay creates a new Relay instance with a specified buffer size
+func NewRelay(buffer uint) *Relay {
+	if buffer == 0 {
+		panic("buffer size must be greater than 0")
+	}
+	comm := make(chan Payload, buffer)
+	return &Relay{comm: comm, C: comm}
+}
+
+// Send sends a message to the channel receiver
+// This is non-blocking and returns an error if the channel buffer is full
+func (r *Relay) Send(ctx context.Context, data any) error {
+	select {
+	case r.comm <- Payload{Ctx: common.FreezeContext(ctx), Data: data}:
+		return nil
+	default:
+		return fmt.Errorf("%w: failed to relay <%T>", errChannelBufferFull, data)
+	}
+}
+
+// Close closes the relay channel
+func (r *Relay) Close() {
+	close(r.comm)
+}

exchange/stream/relay_test.go

@@ -0,0 +1,43 @@
+package stream
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestNewRelay(t *testing.T) {
+	t.Parallel()
+	assert.Panics(t, func() { NewRelay(0) }, "buffer size should be greater than 0")
+	r := NewRelay(5)
+	require.NotNil(t, r)
+	assert.Equal(t, 5, cap(r.comm))
+}
+
+func TestSend(t *testing.T) {
+	t.Parallel()
+	r := NewRelay(1)
+	require.NotNil(t, r)
+	assert.NoError(t, r.Send(t.Context(), "test"))
+	assert.ErrorIs(t, r.Send(t.Context(), "overflow"), errChannelBufferFull)
+}
+
+func TestRead(t *testing.T) {
+	t.Parallel()
+	r := NewRelay(1)
+	require.NotNil(t, r)
+	require.Empty(t, r.C)
+	assert.NoError(t, r.Send(t.Context(), "test"))
+	require.Len(t, r.C, 1)
+	assert.Equal(t, "test", (<-r.C).Data)
+}
+
+func TestClose(t *testing.T) {
+	t.Parallel()
+	r := NewRelay(1)
+	require.NotNil(t, r)
+	r.Close()
+	_, ok := <-r.C
+	assert.False(t, ok)
+}