refact: buffer hop mux incoming channel for better burst handling
Also merges granular yamux stream opening into OpenStream.
Lee Yunjin committed
May 12, 2026 at 10:57 UTC
dceb51e08b66d491fc1d5d4aa7c82f4b7ce5b9a3
1 file changed
+27
-95
portal/overlay/hop_mux.go
+27
-95
@@ -7,7 +7,6 @@ import (
7
"fmt"
8
"io"
9
"net"
10
- "strings"
10
"sync"
11
"time"
12
@@ -46,7 +45,7 @@ func NewHopMux(overlay *Overlay) (*HopMux, error) {
45
return &HopMux{
46
listener: listener,
47
overlay: overlay,
49
- incoming: make(chan HopStream),
48
+ incoming: make(chan HopStream, 1024),
49
outbound: make(map[string]*yamux.Session),
50
done: make(chan struct{}),
51
}, nil
@@ -60,17 +59,13 @@ func (m *HopMux) Serve(ctx context.Context) error {
59
60
for {
61
conn, err := m.listener.Accept()
63
- switch {
64
- case err == nil:
65
- go m.serveSession(ctx, conn)
66
- case errors.Is(err, net.ErrClosed):
67
- return nil
68
- default:
69
- if ctxErr := ctx.Err(); ctxErr != nil {
70
- return ctxErr
62
+ if err != nil {
63
+ if errors.Is(err, net.ErrClosed) || ctx.Err() != nil {
64
+ return ctx.Err()
65
}
66
return fmt.Errorf("accept hop mux connection: %w", err)
67
}
68
+ go m.serveSession(ctx, conn)
69
}
70
}
71
@@ -92,10 +87,7 @@ func (m *HopMux) Close() error {
87
default:
88
}
89
close(m.done)
95
- sessions := make([]*yamux.Session, 0, len(m.outbound))
96
- for _, session := range m.outbound {
97
- sessions = append(sessions, session)
98
- }
90
+ sessions := m.outbound
91
m.outbound = make(map[string]*yamux.Session)
92
m.mu.Unlock()
93
@@ -107,21 +99,15 @@ func (m *HopMux) Close() error {
99
}
100
101
func (m *HopMux) OpenStream(ctx context.Context, overlayIPv4, token string) (net.Conn, error) {
110
- token = strings.TrimSpace(token)
111
- if token == "" {
112
- return nil, errors.New("next hop token is required")
113
- }
114
-
115
- overlayIPv4 = strings.TrimSpace(overlayIPv4)
116
- if overlayIPv4 == "" {
117
- return nil, errors.New("next hop overlay ipv4 is required")
102
+ if token == "" || overlayIPv4 == "" {
103
+ return nil, errors.New("hop token and overlay ipv4 are required")
104
}
105
106
session, err := m.session(ctx, overlayIPv4)
107
if err != nil {
108
return nil, err
109
}
124
- stream, err := openYamuxStream(ctx, session)
110
+ stream, err := session.OpenStream()
111
if err != nil {
112
m.mu.Lock()
113
if m.outbound[overlayIPv4] == session {
@@ -140,50 +126,22 @@ func (m *HopMux) OpenStream(ctx context.Context, overlayIPv4, token string) (net
126
frame := make([]byte, 4+len(payload))
127
binary.BigEndian.PutUint32(frame[:4], uint32(len(payload)))
128
copy(frame[4:], payload)
143
- if n, err := stream.Write(frame); err != nil {
129
+
130
+ _ = stream.SetWriteDeadline(time.Now().Add(defaultTokenTimeout))
131
+ if _, err := stream.Write(frame); err != nil {
132
_ = stream.Close()
133
return nil, err
146
- } else if n != len(frame) {
147
- _ = stream.Close()
148
- return nil, io.ErrShortWrite
134
}
135
+ _ = stream.SetWriteDeadline(time.Time{})
136
return stream, nil
137
}
138
153
-func openYamuxStream(ctx context.Context, session *yamux.Session) (*yamux.Stream, error) {
154
- type openResult struct {
155
- stream *yamux.Stream
156
- err error
157
- }
158
- result := make(chan openResult, 1)
159
- go func() {
160
- stream, err := session.OpenStream()
161
- result <- openResult{stream: stream, err: err}
162
- }()
163
- select {
164
- case res := <-result:
165
- return res.stream, res.err
166
- case <-ctx.Done():
167
- _ = session.Close()
168
- return nil, ctx.Err()
169
- }
170
-}
171
-
139
func (m *HopMux) session(ctx context.Context, overlayIPv4 string) (*yamux.Session, error) {
140
m.mu.Lock()
174
- select {
175
- case <-m.done:
141
+ if session := m.outbound[overlayIPv4]; session != nil && !session.IsClosed() {
142
m.mu.Unlock()
177
- return nil, net.ErrClosed
178
- default:
143
+ return session, nil
144
}
180
- if session := m.outbound[overlayIPv4]; session != nil {
181
- if !session.IsClosed() {
182
- m.mu.Unlock()
183
- return session, nil
184
- }
185
- }
186
- delete(m.outbound, overlayIPv4)
145
m.mu.Unlock()
146
147
addr := net.JoinHostPort(overlayIPv4, fmt.Sprintf("%d", DefaultPeerYamuxPort))
@@ -199,20 +157,12 @@ func (m *HopMux) session(ctx context.Context, overlayIPv4 string) (*yamux.Sessio
157
}
158
159
m.mu.Lock()
202
- select {
203
- case <-m.done:
204
- m.mu.Unlock()
205
- _ = session.Close()
206
- return nil, net.ErrClosed
207
- default:
208
- }
160
+ defer m.mu.Unlock()
161
if current := m.outbound[overlayIPv4]; current != nil && !current.IsClosed() {
210
- m.mu.Unlock()
162
_ = session.Close()
163
return current, nil
164
}
165
m.outbound[overlayIPv4] = session
215
- m.mu.Unlock()
166
return session, nil
167
}
168
@@ -222,30 +172,16 @@ func (m *HopMux) serveSession(ctx context.Context, conn net.Conn) {
172
_ = conn.Close()
173
return
174
}
225
- m.mu.Lock()
226
- select {
227
- case <-m.done:
228
- m.mu.Unlock()
229
- _ = session.Close()
230
- return
231
- default:
232
- }
233
- done := m.done
234
- m.mu.Unlock()
235
- sessionDone := make(chan struct{})
175
+ defer session.Close()
176
+
177
go func() {
178
select {
179
case <-ctx.Done():
180
_ = session.Close()
240
- case <-done:
181
+ case <-m.done:
182
_ = session.Close()
242
- case <-sessionDone:
183
}
184
}()
245
- defer func() {
246
- close(sessionDone)
247
- _ = session.Close()
248
- }()
185
186
for {
187
stream, err := session.AcceptStream()
@@ -275,23 +211,19 @@ func (m *HopMux) handleStream(ctx context.Context, stream *yamux.Stream) {
211
_ = stream.Close()
212
return
213
}
278
- token := strings.TrimSpace(string(payload))
279
- if token == "" {
280
- _ = stream.Close()
281
- return
282
- }
214
+ token := string(payload)
215
+
216
remoteAddr := ""
284
- if stream.RemoteAddr() != nil {
285
- remoteAddr = stream.RemoteAddr().String()
217
+ if addr := stream.RemoteAddr(); addr != nil {
218
+ remoteAddr = addr.String()
219
}
220
+
221
select {
288
- case m.incoming <- HopStream{
289
- Conn: stream,
290
- Token: token,
291
- RemoteAddr: remoteAddr,
292
- }:
222
+ case m.incoming <- HopStream{Conn: stream, Token: token, RemoteAddr: remoteAddr}:
223
case <-ctx.Done():
224
_ = stream.Close()
225
+ case <-m.done:
226
+ _ = stream.Close()
227
}
228
}
229