websocket: Initial refactor for websocket handling to support context propagation

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

exchange/message/message.go

@@ -0,0 +1,47 @@
+package message
+
+import (
+	"context"
+	"errors"
+	"fmt"
+)
+
+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 {
+	// TODO: remove context from payload see: https://github.com/thrasher-corp/gocryptotrader/pull/2066#discussion_r2501403057
+	Ctx  context.Context //nolint:containedctx // context needed for tracing/metrics
+	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: 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/message/message_test.go

@@ -0,0 +1,43 @@
+package message
+
+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)
+}