websocket: Initial refactor for websocket handling to support context propagation #2066
Conversation
- Updated websocket read and handle functions to accept context.Context for better cancellation and timeout management. - Modified test cases to pass context when invoking websocket handling methods. - Enhanced data handler to use a message relay for improved message processing. - Adjusted various functions across the exchange package to ensure context is utilized where necessary.
shazbert
requested review from
cranktakular,
gbjk,
gloriousCode,
samuael and
thrasher-
|
@codex please review |
There was a problem hiding this comment.
Pull Request Overview
This PR implements a significant refactor of the websocket handling system to support context propagation throughout the websocket data flow. The primary motivation is to enable comprehensive metrics and tracing from websocket readers through to external consumers, eliminating circular dependencies and improving the overall architecture.
Key changes include:
- Replaces direct channel usage with a new
message.Relayabstraction for websocket data handling - Adds context parameter propagation through all websocket handling methods
- Removes the
observeDatamethod which was causing overhead by moving data through unnecessary channels
Reviewed Changes
Copilot reviewed 69 out of 69 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| exchanges/trade/trade_types.go | Updates Trade struct to use message.Relay instead of direct channel |
| exchanges/trade/trade.go | Modifies Setup and Update methods to work with message.Relay and context |
| exchanges/sharedtestvalues/sharedtestvalues.go | Updates test websocket creation to use message.Relay |
| exchanges/fill/fill_types.go | Changes Fills struct to use message.Relay for data handling |
| exchanges/fill/fill.go | Updates fill processing methods to support context and message.Relay |
| exchanges/exchange.go | Updates websocket-related method calls to include context parameters |
| Multiple exchange websocket files | Extensive updates to websocket handlers adding context parameters and using message.Relay.Send() |
| } | ||
| oSide, err = order.StringToOrderSide(notification.Data[i].OrderMode) | ||
|
|
||
| oSide, err = order.StringToOrderSide(strings.ReplaceAll(notification.Data[i].OrderMode, "MODE_", "")) |
There was a problem hiding this comment.
[nitpick] This string replacement logic appears to be handling a specific data format transformation. Consider extracting this into a named constant or helper function for better maintainability and to document the purpose of removing the 'MODE_' prefix.
| oSide, err = order.StringToOrderSide(strings.ReplaceAll(notification.Data[i].OrderMode, "MODE_", "")) | |
| oSide, err = order.StringToOrderSide(stripOrderModePrefix(notification.Data[i].OrderMode)) |
There was a problem hiding this comment.
This was just to fix a test, the other implementations don't seem to have the same handling which leads me to believe the test websocket payload might be wrong but this is all out of scope.
| } | ||
| errCode, err := jsonparser.GetInt(respRaw, "code") | ||
| if err != nil { | ||
| log.Errorf(log.ExchangeSys, "%s %s 'code': %s from message: %s", e.Name, common.ErrParsingWSField, err, respRaw) | ||
| } | ||
| return fmt.Errorf("WS auth subscription error; Status: %s Error Code: %d", status, errCode) |
There was a problem hiding this comment.
This code will always return an error even when authentication succeeds. The error return should be conditional - it should only return an error when the status indicates failure, not success.
| } | |
| errCode, err := jsonparser.GetInt(respRaw, "code") | |
| if err != nil { | |
| log.Errorf(log.ExchangeSys, "%s %s 'code': %s from message: %s", e.Name, common.ErrParsingWSField, err, respRaw) | |
| } | |
| return fmt.Errorf("WS auth subscription error; Status: %s Error Code: %d", status, errCode) | |
| } else if status == "ERROR" { | |
| errCode, err := jsonparser.GetInt(respRaw, "code") | |
| if err != nil { | |
| log.Errorf(log.ExchangeSys, "%s %s 'code': %s from message: %s", e.Name, common.ErrParsingWSField, err, respRaw) | |
| } | |
| return fmt.Errorf("WS auth subscription error; Status: %s Error Code: %d", status, errCode) | |
| } | |
| // For any other status, do not return an error, just log it for visibility | |
| log.Warnf(log.ExchangeSys, "%s received unhandled WS auth status: %s from message: %s", e.Name, status, respRaw) | |
| return nil |
There was a problem hiding this comment.
status return might not ever be "ERROR" and this doesn't take in an unhandled status event; as this is out of scope I will defer this.
| for i := range snapBundle { | ||
| if positionData, ok := snapBundle[i].([]any); ok { | ||
| e.wsHandleOrder(positionData) | ||
| return e.wsHandleOrder(ctx, positionData) |
There was a problem hiding this comment.
This return statement breaks out of the loop after processing only the first position. The loop should continue processing all positions and only return after the loop completes or if an error occurs.
| return e.wsHandleOrder(ctx, positionData) | |
| if err := e.wsHandleOrder(ctx, positionData); err != nil { | |
| return err | |
| } |
| ctx := account.DeployCredentialsToContext(t.Context(), &account.Credentials{ | ||
| Key: "testkey", | ||
| Secret: base64.StdEncoding.EncodeToString([]byte("testsecret")), | ||
| }) |
There was a problem hiding this comment.
[nitpick] Consider extracting the test credentials creation into a helper function since this pattern may be reused across multiple test functions. This would improve maintainability and reduce code duplication.
There was a problem hiding this comment.
Unless others want me to do this I will leave this open.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
Codex Review: Here are some suggestions.
About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you open a pull request for review, mark a draft as ready, or comment "@codex review". If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex fix this CI failure" or "@codex address that feedback".
exchange/websocket/buffer/buffer.go
Outdated
|
|
||
| // Publish all state changes, disregarding verbosity or sync requirements. | ||
| holder.ob.Publish() | ||
| o.dataHandler <- holder.ob | ||
| o.dataHandler.Send(ctx, holder.ob) |
There was a problem hiding this comment.
[P1] Handle relay Send errors to avoid dropping websocket data
The new message.Relay reports a full buffer via an error, but both Orderbook.LoadSnapshot and Orderbook.Update call o.dataHandler.Send(ctx, holder.ob) and ignore the return value. When the buffer is full these sends will now fail silently, so consumers never see the published order book state even though internal structures were updated. Previously the channel send blocked until consumers drained the queue, so no data was lost. The same pattern also exists in exchanges/trade/trade.go and exchanges/fill/fill.go. Consider propagating or at least logging the error so backpressure can be surfaced and the update retried or dropped explicitly.
Useful? React with 👍 / 👎.
engine/rpcserver.go
Outdated
| return nil, fmt.Errorf("websocket not supported for exchange %s", r.Exchange) | ||
| } | ||
|
|
||
| err = w.SetProxyAddress(r.Proxy) | ||
| err = w.SetProxyAddress(ctx, r.Proxy) | ||
| if err != nil { |
There was a problem hiding this comment.
[P1] Do not reuse RPC request context for websocket lifetime
The RPC handlers now pass the request’s ctx into w.Enable and w.SetProxyAddress. Connect stores this context and uses it for future reconnects inside the connection monitor. gRPC cancels the request context as soon as the RPC returns, so any later reconnect attempt will run with a canceled context and fail immediately with context canceled, leaving the websocket permanently disconnected after the first dropout or proxy change. These long‑lived operations should use a background context or derive a new context instead of the short‑lived RPC context.
Useful? React with 👍 / 👎.
shazbert
added
reconstructing
|
@codex please review |
|
Codex Review: Didn't find any major issues. What shall we delve into next? About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you open a pull request for review, mark a draft as ready, or comment "@codex review". If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex fix this CI failure" or "@codex address that feedback". |
shazbert
added
review me
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #2066 +/- ##
==========================================
- Coverage 42.00% 41.91% -0.10%
==========================================
Files 443 444 +1
Lines 143107 142800 -307
==========================================
- Hits 60111 59850 -261
+ Misses 75845 75773 -72
- Partials 7151 7177 +26
🚀 New features to boost your workflow:
|
gloriousCode
left a comment
gloriousCode
left a comment
There was a problem hiding this comment.
I appreciate context usage, but the nolint leaves me uncertain. I will leave the AI output for you to consider
gbjk
left a comment
gbjk
left a comment
There was a problem hiding this comment.
I'd like to avoid passing the context down the channel, as per this thread
gbjk
left a comment
gbjk
left a comment
There was a problem hiding this comment.
Great work, thanks.
Mini-review.
I'm also thinking that the justification for exchange/message being a package might be solvable by using the consumer defined interface pattern. I'll look into that and report back
common/common.go
Outdated
| for k := range fc.values { | ||
| if ctx.Value(k) != nil { | ||
| return nil, fmt.Errorf("%w: %q", errDuplicateContextKey, k) | ||
| } | ||
| } | ||
| return &mergeCtx{ctx, 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.values[key]; ok { | ||
| return val | ||
| } | ||
| return m.Context.Value(key) | ||
| } |
There was a problem hiding this comment.
🚧 Change request:
I think this code was maybe a reaction to the fatcontext linter.
The idiomatic way to set keys on contexts is with context.WithValue.
It's fine that they're nested, that's the way go contexts were designed.
The linter's author's use-case was to not call context.WithContext N+ times in a for loop and keep overwriting the same key.
They highlighted this in the blog post.
They should have been using a shadow var, and that's the solution they expect to be used for this linter warning.
None of that is relevant to us.
We're looping over a list of values and going to return a context with all of them.
Calling WithValue is the right thing to do.
| for k := range fc.values { | |
| if ctx.Value(k) != nil { | |
| return nil, fmt.Errorf("%w: %q", errDuplicateContextKey, k) | |
| } | |
| } | |
| return &mergeCtx{ctx, 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.values[key]; ok { | |
| return val | |
| } | |
| return m.Context.Value(key) | |
| } | |
| // MergeCtx adds the frozen values to an existing context | |
| func MergeCtx(ctx context.Context, fc FrozenContext) (context.Context, error) { | |
| for k, v := range fc.values { | |
| if ctx.Value(k) != nil { | |
| return nil, fmt.Errorf("%w: %q", errDuplicateContextKey, k) | |
| } | |
| ctx = context.WithValue(ctx, k, v) //nolint:fatcontext // false-positive; linter catches non-shadow contexts, but we return this ctx | |
| } | |
| return ctx, nil | |
| } |
There was a problem hiding this comment.
I benched it and it indeed nests contexts within each other even with different key values, so the implementation I have done can stay:
// 100000000 12.17 ns/op 0 B/op 0 allocs/op (with mergeCtx)
//
// 38971 26070 ns/op 0 B/op 0 allocs/op (using context.WithValue)
func BenchmarkXxx(b *testing.B) {
m := FrozenContext{
values: make(map[any]any),
}
for i := range 10000 {
m.values[i] = i
}
ctx, err := ThawCtx(m)
require.NoError(b, err)
b.ResetTimer()
for b.Loop() {
if val := ctx.Value(5000); val == nil {
b.Fatal("value not found")
}
}
}Can you please check yourself to see if I have done something silly @gbjk
There was a problem hiding this comment.
@shazbert Whoa. Yeah, you've benched something nobody should ever do, because you've stored thousands of keys.
I wasn't saying nesting contexts is a bad thing.
I was saying nesting thousands of contexts like the author did is a bad thing.
If you're in a situation where you're storing a large number of keys then storing a single struct would make sense.
But we're not. If we store up to 5 keys, then accessing them shouldn't take longer.
See this bench
Accessing these keys is about the same time (faster for 4, slower for 10).
However, plot twist
Freezing the context takes N* longer, though.
And our use-case freezes the context a lot more than normal, since we're doing it for every message we want to relay, right?
So If we're storing even 2 keys, all the time, then we should use your method.
If you agree, just resolve.
There was a problem hiding this comment.
Leave it up to @thrasher- @gloriousCode they can resolve or whatnot
Id opt to stay as is so that if somebody has an ungodly amount (I don't know why) of things for context at least its collapsed at this point 🤷.
Also I can do what metadata does and have a single key lookup for the map on a context and use context.WithValue() so we don't need to do the nolint on a struct as an option, if we want to keep the storage of the map?
There was a problem hiding this comment.
I think it's fine as is. No need to use a real context.WithValue.
I'd lean towards saying: Use this context thing from the start, though.
Make it an app context that implements context.Context, backed by context.Context for lifecycle management but with a different value store permanently attached to WithValue and Value for speed.
Because request cycle values are a thing, a context is the right place to store them, and native context.* is bad at it.
The nice thing about that is that nested context for any other reason, like cancellation, would still un-nest to find your context's value store. And you can still delegate to the underlying context Value like you do.
There was a problem hiding this comment.
Thank you both for looking into solutions so thoroughly! I was only after a message package specific helper to grab OTELs metrics. I find the current committed solution addresses the problem I raised. I hope its not too complex on the Thawed side, Shazbert.
I think there's room for some improvements to be made, but design nicety improvements should be their own PR. I think the only place to try and improve upon now is efficiency. This looks like it could have an impact since it occurs for every websocket message. I'll look into that today and leave review comments
gbjk
left a comment
gbjk
left a comment
There was a problem hiding this comment.
Two small change requests that are unimportant.
https://github.com/gbjk/gocryptotrader/commit/3506efa04.patch
I was looking at the option of what we discussed, and saw advantages to moving to contextStore type.
The issue is that it becomes mutable.
I haven't finished this, but it's 90% of the way there.
I'll just throw this ball to you and you can use it, drop it, whatever.
The main realisation I had was that we're paying the cost of WithValue already N+ times with the current design.
Yes, we avoid it at Thaw. But we paid it originally putting values into the context on the other side before Freeze.
Co-authored-by: Gareth Kirwan <[email protected]>
Co-authored-by: Gareth Kirwan <[email protected]>
Yeah I like it. All the values that are put in by the library are captured across this boundary by freezecontext so you don't need to set the required values you want in common, it will auto scale. Anything else are dropped and it keeps its speed. Will see if @thrasher- or @gloriousCode wants this for this PR. Then I will drop it in. Or we can open another PR for it after. |
gloriousCode
left a comment
gloriousCode
left a comment
There was a problem hiding this comment.
I have attempted to improve things. I haven't come up short, but I haven't smashed face either
Happy as is, it works. Can always adjust it once the patch is more fleshed out and in a different PR |
gbjk
left a comment
gbjk
left a comment
There was a problem hiding this comment.
Good work.
Couple of nits and some pedantry about package naming.
Otherwise I think I'll be 👍 on the next round.
engine/rpcserver.go
Outdated
| // 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 { |
There was a problem hiding this comment.
🐒 Optional nitpick:
Okay; I kinda get it, except LongLivedSession seems ... well weird.
We're basically just stripping out the lifecycle from the context, right?
So ... isn't this just a very fancy context.WithoutCancel ?
There was a problem hiding this comment.
Yeah good point, any other naming suggestions and I will change it, and should this function live somewhere like common?
There was a problem hiding this comment.
I mean... if it's context.WithoutCancel you don't need this function... right?
Just drop in replace it.
There was a problem hiding this comment.
Well I feel like an idiot. 😆
exchange/message/message.go
Outdated
| @@ -0,0 +1,48 @@ | |||
| package message | |||
There was a problem hiding this comment.
🚧 Change request:
exchange/message feels like a very overloaded term and it feels like you're domain-squatting it a bit here.
I get that this could represent any message queue.
But right now we don't have an idea of any other type of messages here.
And this is fundamentally a relay, not a message thing.
So I'd like to see this in exchange/websocket/relay/relay.go instead.
But I'm also thinking it's even better in exchange/websocket/buffer/relay.go 🎉
There was a problem hiding this comment.
My reasoning for placing this in exchange/message was that the relay is intentionally generic — it doesn’t depend on WebSockets or exchange code, and its API (message.Relay) is transport-agnostic by design.
If we move it into websocket/relay, it implicitly couples the type to a subsystem it doesn’t rely on and might make it harder to reuse later for other messaging paths (e.g., FIX, -- these are all a strech --> RPC, internal event pumps, REST batch queues).
That said, I’m happy to relocate it if the consensus is that the Relay should remain WebSocket-specific for now.
Also Buffer naming I think will be subject to change because no websocket orderbooks should buffer events or sort them.
Also what about exchange/stream/relay.go?
There was a problem hiding this comment.
I think exchange/stream/relay.go is the best. Removes the websocket coupling since it's protocol agnostic and I agree that message is too generic
There was a problem hiding this comment.
exchange/stream/relay.gois good 👍
🥃 Your reasoning about buffer was backwards. Just because websocket orderbooks are in a package called buffer doesn't mean that nothing else can be, or that when they move out we couldn't leave a relay in there. But that's just an aside.
There was a problem hiding this comment.
I was more meaning, its going to change, then we will have to change the package name for all relay instances and it doesn't really reflect buffer that's all.
Co-authored-by: Gareth Kirwan <[email protected]>
gbjk
left a comment
gbjk
left a comment
There was a problem hiding this comment.
Good work.
One change request, but optional.
4 participants
Adds package
messagefor the websocket data handler so there is no circular dependenciesSendSendtakes in context that can be consumed for future metrics/trace. e.g. external consumer using signal (orderbook change) can get metrics all the way back from the websocket reader routine.This removes
observeDatamethod on websocket manager. This was just moving data to another channel just to indicate blockage which had a bit of overhead, the above design eliminates the need for that mechanism.Tries to keep things contained in websocket only a few context.TODO's and these will be opened up in a future PR:
This complements issue #2026 for a context based OTEL solution
Also complements issue #2068
This also will help trace along with profiling #1951 as it seems to be just a garbage collection spike, optionally can track GC cycles on full trace for that signal.
Type of change
Please delete options that are not relevant and add an
xin[]as item is complete.How has this been tested
Please describe the tests that you ran to verify your changes. Provide instructions so we can reproduce. Please also list any relevant details for your test configuration and
also consider improving test coverage whilst working on a certain feature or package.
Checklist