jay/go-honeybee
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Injected context cancellation for dial and retry cancellation.

Browse Source
This commit is contained in:
Jay
2026-04-18 17:11:22 -04:00
parent e1cdc1cf9c
commit b4c5c897e8
12 changed files with 182 additions and 230 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
initiatorpool/config.go
View File

@@ -1,13 +1,14 @@
package initiatorpool
import (
"context"
"git.wisehodl.dev/jay/go-honeybee/transport"
"time"
)
// Types
type WorkerFactory func(id string, stop <-chan struct{}) (*Worker, error)
type WorkerFactory func(ctx context.Context, id string) (*Worker, error)
// Pool Config

38
initiatorpool/pool.go
View File

@@ -1,6 +1,7 @@
package initiatorpool
import (
"context"
"git.wisehodl.dev/jay/go-honeybee/transport"
"git.wisehodl.dev/jay/go-honeybee/types"
"log/slog"
@@ -13,14 +14,12 @@ import (
type Peer struct {
id string
worker *Worker
stop chan struct{}
}
type WorkerContext struct {
Inbox chan<- InboxMessage
Events chan<- PoolEvent
Errors chan<- error
PoolDone <-chan struct{}
Logger *slog.Logger
Dialer types.Dialer
ConnectionConfig *transport.ConnectionConfig
@@ -47,11 +46,13 @@ type PoolEvent struct {
// Pool
type Pool struct {
ctx context.Context
cancel context.CancelFunc
peers map[string]*Peer
inbox chan InboxMessage
events chan PoolEvent
errors chan error
done chan struct{}
dialer types.Dialer
config *PoolConfig
@@ -62,7 +63,8 @@ type Pool struct {
closed bool
}
func NewPool(config *PoolConfig, logger *slog.Logger) (*Pool, error) {
func NewPool(ctx context.Context, config *PoolConfig, logger *slog.Logger,
) (*Pool, error) {
if config == nil {
config = GetDefaultPoolConfig()
}
@@ -71,8 +73,9 @@ func NewPool(config *PoolConfig, logger *slog.Logger) (*Pool, error) {
// The factory function should be non-blocking or else Connect() may cause
// deadlocks.
if config.WorkerFactory == nil {
config.WorkerFactory = func(id string, stop <-chan struct{}) (*Worker, error) {
return NewWorker(id, stop, config.WorkerConfig)
config.WorkerFactory = func(
ctx context.Context, id string) (*Worker, error) {
return NewWorker(ctx, id, config.WorkerConfig)
}
}
@@ -80,12 +83,15 @@ func NewPool(config *PoolConfig, logger *slog.Logger) (*Pool, error) {
return nil, err
}
pctx, cancel := context.WithCancel(ctx)
p := &Pool{
ctx: pctx,
cancel: cancel,
peers: make(map[string]*Peer),
inbox: make(chan InboxMessage, 256),
events: make(chan PoolEvent, 10),
errors: make(chan error, 10),
done: make(chan struct{}),
dialer: transport.NewDialer(),
config: config,
logger: logger,
@@ -125,17 +131,12 @@ func (p *Pool) Close() {
}
p.closed = true
close(p.done)
p.cancel()
peers := p.peers
p.peers = make(map[string]*Peer)
p.mu.Unlock()
for _, p := range peers {
close(p.stop)
}
go func() {
p.wg.Wait()
close(p.inbox)
@@ -162,13 +163,9 @@ func (p *Pool) Connect(id string) error {
return NewPoolError("connection already exists")
}
// Create new worker
stop := make(chan struct{})
// The worker factory must be non-blocking to avoid deadlocks
worker, err := p.config.WorkerFactory(id, stop)
worker, err := p.config.WorkerFactory(p.ctx, id)
if err != nil {
close(stop)
return err
}
@@ -180,7 +177,6 @@ func (p *Pool) Connect(id string) error {
Inbox: p.inbox,
Events: p.events,
Errors: p.errors,
PoolDone: p.done,
Logger: logger,
Dialer: p.dialer,
ConnectionConfig: p.config.ConnectionConfig,
@@ -189,7 +185,7 @@ func (p *Pool) Connect(id string) error {
p.wg.Add(1)
go worker.Start(ctx, &p.wg)
p.peers[id] = &Peer{id: id, worker: worker, stop: stop}
p.peers[id] = &Peer{id: id, worker: worker}
return nil
}
@@ -214,7 +210,7 @@ func (p *Pool) Remove(id string) error {
delete(p.peers, id)
p.mu.Unlock()
close(peer.stop)
peer.worker.Stop()
return nil
}

28
initiatorpool/pool_test.go
View File

@@ -1,6 +1,7 @@
package initiatorpool
import (
"context"
"fmt"
"git.wisehodl.dev/jay/go-honeybee/honeybeetest"
"git.wisehodl.dev/jay/go-honeybee/transport"
@@ -17,12 +18,12 @@ func _TestPoolConnect(t *testing.T) {
t.Run("successfully adds connection", func(t *testing.T) {
mockSocket := honeybeetest.NewMockSocket()
mockDialer := &honeybeetest.MockDialer{
DialFunc: func(string, http.Header) (types.Socket, *http.Response, error) {
DialContextFunc: func(context.Context, string, http.Header) (types.Socket, *http.Response, error) {
return mockSocket, nil, nil
},
}
pool, err := NewPool(nil, nil)
pool, err := NewPool(context.Background(), nil, nil)
assert.NoError(t, err)
pool.dialer = mockDialer
@@ -48,12 +49,12 @@ func _TestPoolConnect(t *testing.T) {
t.Run("does not add duplicate", func(t *testing.T) {
mockSocket := honeybeetest.NewMockSocket()
mockDialer := &honeybeetest.MockDialer{
DialFunc: func(string, http.Header) (types.Socket, *http.Response, error) {
DialContextFunc: func(context.Context, string, http.Header) (types.Socket, *http.Response, error) {
return mockSocket, nil, nil
},
}
pool, err := NewPool(nil, nil)
pool, err := NewPool(context.Background(), nil, nil)
assert.NoError(t, err)
pool.dialer = mockDialer
@@ -74,6 +75,7 @@ func _TestPoolConnect(t *testing.T) {
t.Run("fails to add connection", func(t *testing.T) {
pool, err := NewPool(
context.Background(),
&PoolConfig{
ConnectionConfig: &transport.ConnectionConfig{
Retry: &transport.RetryConfig{
@@ -84,7 +86,7 @@ func _TestPoolConnect(t *testing.T) {
}, nil)
assert.NoError(t, err)
pool.dialer = &honeybeetest.MockDialer{
DialFunc: func(string, http.Header) (types.Socket, *http.Response, error) {
DialContextFunc: func(context.Context, string, http.Header) (types.Socket, *http.Response, error) {
return nil, nil, fmt.Errorf("dial failed")
},
}
@@ -111,12 +113,12 @@ func _TestPoolRemove(t *testing.T) {
t.Run("removes known url", func(t *testing.T) {
mockSocket := honeybeetest.NewMockSocket()
mockDialer := &honeybeetest.MockDialer{
DialFunc: func(string, http.Header) (types.Socket, *http.Response, error) {
DialContextFunc: func(context.Context, string, http.Header) (types.Socket, *http.Response, error) {
return mockSocket, nil, nil
},
}
pool, err := NewPool(nil, nil)
pool, err := NewPool(context.Background(), nil, nil)
assert.NoError(t, err)
pool.dialer = mockDialer
@@ -139,12 +141,12 @@ func _TestPoolRemove(t *testing.T) {
t.Run("unknown url returns error", func(t *testing.T) {
mockSocket := honeybeetest.NewMockSocket()
mockDialer := &honeybeetest.MockDialer{
DialFunc: func(string, http.Header) (types.Socket, *http.Response, error) {
DialContextFunc: func(context.Context, string, http.Header) (types.Socket, *http.Response, error) {
return mockSocket, nil, nil
},
}
pool, err := NewPool(nil, nil)
pool, err := NewPool(context.Background(), nil, nil)
assert.NoError(t, err)
pool.dialer = mockDialer
@@ -156,12 +158,12 @@ func _TestPoolRemove(t *testing.T) {
t.Run("closed pool returns error", func(t *testing.T) {
mockSocket := honeybeetest.NewMockSocket()
mockDialer := &honeybeetest.MockDialer{
DialFunc: func(string, http.Header) (types.Socket, *http.Response, error) {
DialContextFunc: func(context.Context, string, http.Header) (types.Socket, *http.Response, error) {
return mockSocket, nil, nil
},
}
pool, err := NewPool(nil, nil)
pool, err := NewPool(context.Background(), nil, nil)
assert.NoError(t, err)
pool.dialer = mockDialer
@@ -184,12 +186,12 @@ func _TestPoolSend(t *testing.T) {
return nil
}
mockDialer := &honeybeetest.MockDialer{
DialFunc: func(string, http.Header) (types.Socket, *http.Response, error) {
DialContextFunc: func(context.Context, string, http.Header) (types.Socket, *http.Response, error) {
return mockSocket, nil, nil
},
}
pool, err := NewPool(nil, nil)
pool, err := NewPool(context.Background(), nil, nil)
assert.NoError(t, err)
pool.dialer = mockDialer

79
initiatorpool/worker.go
View File

@@ -2,6 +2,7 @@ package initiatorpool
import (
"container/list"
"context"
"git.wisehodl.dev/jay/go-honeybee/transport"
"sync"
"time"
@@ -15,15 +16,16 @@ type receivedMessage struct {
}
type Worker struct {
ctx context.Context
cancel context.CancelFunc
id string
stop <-chan struct{}
config *WorkerConfig
outbound chan []byte
}
func NewWorker(
ctx context.Context,
id string,
stop <-chan struct{},
config *WorkerConfig,
) (*Worker, error) {
@@ -36,9 +38,11 @@ func NewWorker(
return nil, err
}
wctx, cancel := context.WithCancel(ctx)
w := &Worker{
ctx: wctx,
cancel: cancel,
id: id,
stop: stop,
outbound: make(chan []byte, 64),
config: config,
}
@@ -50,7 +54,7 @@ func (w *Worker) Send(data []byte) error {
select {
case w.outbound <- data:
return nil
case <-w.stop:
case <-w.ctx.Done():
return NewWorkerError(w.id, "worker is stopped")
default:
return NewWorkerError(w.id, "outbound queue full")
@@ -63,7 +67,14 @@ func (w *Worker) Start(
) {
}
func (w *Worker) Stop() {
w.cancel()
}
func (w *Worker) runSession(
ctx context.Context,
wctx WorkerContext,
messages chan<- receivedMessage,
heartbeat chan<- struct{},
dial chan<- struct{},
@@ -71,11 +82,6 @@ func (w *Worker) runSession(
keepalive <-chan struct{},
outbound <-chan []byte,
newConn <-chan *transport.Connection,
ctx WorkerContext,
workerStop <-chan struct{},
poolDone <-chan struct{},
) {
}
@@ -98,20 +104,18 @@ func (w *Worker) runWriter(
}
func (w *Worker) runStopMonitor(
ctx context.Context,
conn *transport.Connection,
keepalive <-chan struct{},
workerStop <-chan struct{},
poolDone <-chan struct{},
sessionDone <-chan struct{},
onStop func(),
) {
}
func (w *Worker) runForwarder(
ctx context.Context,
messages <-chan receivedMessage,
inbox chan<- InboxMessage,
stop <-chan struct{},
poolDone <-chan struct{},
maxQueueSize int,
) {
queue := list.New()
@@ -129,9 +133,7 @@ func (w *Worker) runForwarder(
}
select {
case <-stop:
return
case <-poolDone:
case <-ctx.Done():
return
case msg := <-messages:
// limit queue size if maximum is configured
@@ -154,17 +156,15 @@ func (w *Worker) runForwarder(
}
func (w *Worker) runKeepalive(
ctx context.Context,
heartbeat <-chan struct{},
keepalive chan<- struct{},
stop <-chan struct{},
poolDone <-chan struct{},
) {
// disable keepalive timeout if not configured
if w.config.KeepaliveTimeout <= 0 {
// wait for stop signal and exit
// wait for cancel and exit
select {
case <-stop:
case <-poolDone:
case <-ctx.Done():
}
return
}
@@ -174,9 +174,7 @@ func (w *Worker) runKeepalive(
for {
select {
case <-stop:
return
case <-poolDone:
case <-ctx.Done():
return
case <-heartbeat:
// drain the timer channel and reset
@@ -199,28 +197,29 @@ func (w *Worker) runKeepalive(
}
}
func (w *Worker) dial(ctx WorkerContext) (*transport.Connection, error) {
conn, err := transport.NewConnection(w.id, ctx.ConnectionConfig, ctx.Logger)
func (w *Worker) dial(
ctx context.Context,
wctx WorkerContext,
) (*transport.Connection, error) {
conn, err := transport.NewConnection(w.id, wctx.ConnectionConfig, wctx.Logger)
if err != nil {
return nil, err
}
conn.SetDialer(ctx.Dialer)
return conn, conn.Connect()
conn.SetDialer(wctx.Dialer)
return conn, conn.Connect(ctx)
}
func (w *Worker) runDialer(
ctx context.Context,
wctx WorkerContext,
dial <-chan struct{},
newConn chan<- *transport.Connection,
ctx WorkerContext,
stop <-chan struct{},
poolDone <-chan struct{},
) {
for {
select {
case <-stop:
return
case <-poolDone:
case <-ctx.Done():
return
case <-dial:
// drain dial signals while connection is being established
@@ -236,15 +235,14 @@ func (w *Worker) runDialer(
}()
// dial a new connection
conn, err := w.dial(ctx)
conn, err := w.dial(ctx, wctx)
close(done)
// send error if dial failed and continue
if err != nil {
select {
case ctx.Errors <- err:
case <-stop:
case <-poolDone:
case wctx.Errors <- err:
case <-ctx.Done():
}
continue
}
@@ -252,10 +250,7 @@ func (w *Worker) runDialer(
// send the new connection or close and exit
select {
case newConn <- conn:
case <-stop:
conn.Close()
return
case <-poolDone:
case <-ctx.Done():
conn.Close()
return
}

157
initiatorpool/worker_test.go
View File

@@ -1,6 +1,7 @@
package initiatorpool
import (
"context"
"fmt"
"git.wisehodl.dev/jay/go-honeybee/honeybeetest"
"git.wisehodl.dev/jay/go-honeybee/transport"
@@ -18,11 +19,11 @@ func TestRunForwarder(t *testing.T) {
t.Run("message passes through to inbox", func(t *testing.T) {
messages := make(chan receivedMessage, 1)
inbox := make(chan InboxMessage, 1)
stop := make(chan struct{})
defer close(stop)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
w := &Worker{id: "wss://test"}
go w.runForwarder(messages, inbox, stop, nil, 0)
go w.runForwarder(ctx, messages, inbox, 0)
messages <- receivedMessage{data: []byte("hello"), receivedAt: time.Now()}
@@ -39,8 +40,8 @@ func TestRunForwarder(t *testing.T) {
t.Run("oldest message dropped when queue is full", func(t *testing.T) {
messages := make(chan receivedMessage, 1)
inbox := make(chan InboxMessage, 1)
stop := make(chan struct{})
defer close(stop)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
gate := make(chan struct{})
gatedInbox := make(chan InboxMessage)
@@ -54,7 +55,7 @@ func TestRunForwarder(t *testing.T) {
}()
w := &Worker{id: "wss://test"}
go w.runForwarder(messages, gatedInbox, stop, nil, 2)
go w.runForwarder(ctx, messages, gatedInbox, 2)
// send three messages while the gated inbox is blocked
messages <- receivedMessage{data: []byte("first"), receivedAt: time.Now()}
@@ -83,42 +84,20 @@ func TestRunForwarder(t *testing.T) {
})
t.Run("exits on stop", func(t *testing.T) {
t.Run("exits on context cancellation", func(t *testing.T) {
messages := make(chan receivedMessage, 1)
inbox := make(chan InboxMessage, 1)
stop := make(chan struct{})
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
w := &Worker{id: "wss://test"}
done := make(chan struct{})
go func() {
w.runForwarder(messages, inbox, stop, nil, 0)
w.runForwarder(ctx, messages, inbox, 0)
close(done)
}()
close(stop)
assert.Eventually(t, func() bool {
select {
case <-done:
return true
default:
return false
}
}, honeybeetest.TestTimeout, honeybeetest.TestTick)
})
t.Run("exits on pool done", func(t *testing.T) {
messages := make(chan receivedMessage, 1)
inbox := make(chan InboxMessage, 1)
poolDone := make(chan struct{})
w := &Worker{id: "wss://test"}
done := make(chan struct{})
go func() {
w.runForwarder(messages, inbox, nil, poolDone, 0)
close(done)
}()
close(poolDone)
cancel()
assert.Eventually(t, func() bool {
select {
case <-done:
@@ -134,11 +113,11 @@ func TestRunKeepalive(t *testing.T) {
t.Run("heartbeat resets timer, no keepalive signal fired", func(t *testing.T) {
heartbeat := make(chan struct{}, 3)
keepalive := make(chan struct{}, 1)
stop := make(chan struct{})
defer close(stop)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
w := &Worker{config: &WorkerConfig{KeepaliveTimeout: 100 * time.Millisecond}}
go w.runKeepalive(heartbeat, keepalive, stop, nil)
go w.runKeepalive(ctx, heartbeat, keepalive)
// send heartbeats faster than the timeout
for i := 0; i < 5; i++ {
@@ -160,11 +139,11 @@ func TestRunKeepalive(t *testing.T) {
t.Run("keepalive timeout fires signal", func(t *testing.T) {
heartbeat := make(chan struct{})
keepalive := make(chan struct{}, 1)
stop := make(chan struct{})
defer close(stop)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
w := &Worker{config: &WorkerConfig{KeepaliveTimeout: 20 * time.Millisecond}}
go w.runKeepalive(heartbeat, keepalive, stop, nil)
go w.runKeepalive(ctx, heartbeat, keepalive)
// send no heartbeats, wait for timeout and keepalive signal
assert.Eventually(t, func() bool {
@@ -177,42 +156,19 @@ func TestRunKeepalive(t *testing.T) {
}, honeybeetest.TestTimeout, honeybeetest.TestTick)
})
t.Run("exits on stop", func(t *testing.T) {
t.Run("exits on context cancellation", func(t *testing.T) {
heartbeat := make(chan struct{})
keepalive := make(chan struct{}, 1)
stop := make(chan struct{})
ctx, cancel := context.WithCancel(context.Background())
w := &Worker{config: &WorkerConfig{KeepaliveTimeout: 20 * time.Second}}
done := make(chan struct{})
go func() {
w.runKeepalive(heartbeat, keepalive, stop, nil)
w.runKeepalive(ctx, heartbeat, keepalive)
close(done)
}()
close(stop)
assert.Eventually(t, func() bool {
select {
case <-done:
return true
default:
return false
}
}, honeybeetest.TestTimeout, honeybeetest.TestTick)
})
t.Run("exits on stop", func(t *testing.T) {
heartbeat := make(chan struct{})
keepalive := make(chan struct{}, 1)
poolDone := make(chan struct{})
w := &Worker{config: &WorkerConfig{KeepaliveTimeout: 20 * time.Second}}
done := make(chan struct{})
go func() {
w.runKeepalive(heartbeat, keepalive, nil, poolDone)
close(done)
}()
close(poolDone)
cancel()
assert.Eventually(t, func() bool {
select {
case <-done:
@@ -229,20 +185,20 @@ func TestRunDialer(t *testing.T) {
w := &Worker{id: "wss://test"}
dial := make(chan struct{}, 1)
newConn := make(chan *transport.Connection, 1)
stop := make(chan struct{})
defer close(stop)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
mockSocket := honeybeetest.NewMockSocket()
ctx := WorkerContext{
wctx := WorkerContext{
Errors: make(chan error, 1),
Dialer: &honeybeetest.MockDialer{
DialFunc: func(string, http.Header) (types.Socket, *http.Response, error) {
DialContextFunc: func(context.Context, string, http.Header) (types.Socket, *http.Response, error) {
return mockSocket, nil, nil
},
},
}
go w.runDialer(dial, newConn, ctx, stop, nil)
go w.runDialer(ctx, wctx, dial, newConn)
dial <- struct{}{}
assert.Eventually(t, func() bool {
@@ -260,18 +216,18 @@ func TestRunDialer(t *testing.T) {
w := &Worker{id: "wss://test"}
dial := make(chan struct{}, 1)
newConn := make(chan *transport.Connection, 1)
stop := make(chan struct{})
defer close(stop)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
gate := make(chan struct{})
dialCount := atomic.Int32{}
mockSocket := honeybeetest.NewMockSocket()
connConfig := &transport.ConnectionConfig{Retry: nil} // disable retry
ctx := WorkerContext{
wctx := WorkerContext{
Errors: make(chan error, 1),
Dialer: &honeybeetest.MockDialer{
DialFunc: func(string, http.Header) (types.Socket, *http.Response, error) {
DialContextFunc: func(context.Context, string, http.Header) (types.Socket, *http.Response, error) {
dialCount.Add(1)
<-gate
return mockSocket, nil, nil
@@ -280,7 +236,7 @@ func TestRunDialer(t *testing.T) {
ConnectionConfig: connConfig,
}
go w.runDialer(dial, newConn, ctx, stop, nil)
go w.runDialer(ctx, wctx, dial, newConn)
dial <- struct{}{}
// wait for dial to start blocking on gate
@@ -322,17 +278,19 @@ func TestRunDialer(t *testing.T) {
errors := make(chan error, 1)
dial := make(chan struct{}, 1)
newConn := make(chan *transport.Connection, 1)
stop := make(chan struct{})
defer close(stop)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
// use atomic counter to fail first dial and pass second
dialCount := atomic.Int32{}
mockSocket := honeybeetest.NewMockSocket()
connConfig := &transport.ConnectionConfig{Retry: nil} // disable retry
ctx := WorkerContext{
wctx := WorkerContext{
Errors: errors,
Dialer: &honeybeetest.MockDialer{
DialFunc: func(string, http.Header) (types.Socket, *http.Response, error) {
DialContextFunc: func(
context.Context, string, http.Header,
) (types.Socket, *http.Response, error) {
if dialCount.Add(1) == 1 {
// fail first
return nil, nil, fmt.Errorf("dial failed")
@@ -344,7 +302,7 @@ func TestRunDialer(t *testing.T) {
ConnectionConfig: connConfig,
}
go w.runDialer(dial, newConn, ctx, stop, nil)
go w.runDialer(ctx, wctx, dial, newConn)
dial <- struct{}{}
assert.Eventually(t, func() bool {
@@ -368,21 +326,21 @@ func TestRunDialer(t *testing.T) {
}, honeybeetest.TestTimeout, honeybeetest.TestTick)
})
t.Run("exits on stop", func(t *testing.T) {
t.Run("exits on context cancellation", func(t *testing.T) {
w := &Worker{id: "wss://test"}
dial := make(chan struct{}, 1)
newConn := make(chan *transport.Connection, 1)
stop := make(chan struct{})
ctx, cancel := context.WithCancel(context.Background())
ctx := WorkerContext{Errors: make(chan error, 1)}
wctx := WorkerContext{Errors: make(chan error, 1)}
done := make(chan struct{})
go func() {
w.runDialer(dial, newConn, ctx, stop, nil)
w.runDialer(ctx, wctx, dial, newConn)
close(done)
}()
close(stop)
cancel()
assert.Eventually(t, func() bool {
select {
@@ -393,31 +351,4 @@ func TestRunDialer(t *testing.T) {
}
}, honeybeetest.TestTimeout, honeybeetest.TestTick)
})
t.Run("exits on pool done", func(t *testing.T) {
w := &Worker{id: "wss://test"}
dial := make(chan struct{}, 1)
newConn := make(chan *transport.Connection, 1)
poolDone := make(chan struct{})
ctx := WorkerContext{Errors: make(chan error, 1)}
done := make(chan struct{})
go func() {
w.runDialer(dial, newConn, ctx, nil, poolDone)
close(done)
}()
close(poolDone)
assert.Eventually(t, func() bool {
select {
case <-done:
return true
default:
return false
}
}, honeybeetest.TestTimeout, honeybeetest.TestTick)
})
}