Swap to implementation pattern, make fields and methods public for extensibility.

This commit is contained in:
Jay
2026-04-19 12:49:49 -04:00
parent 2f7e606064
commit d2528d3ac7
10 changed files with 212 additions and 206 deletions
+1 -1
View File
@@ -8,7 +8,7 @@ import (
// Types // Types
type WorkerFactory func(ctx context.Context, id string) (*Worker, error) type WorkerFactory func(ctx context.Context, id string) (Worker, error)
// Pool Config // Pool Config
+2 -2
View File
@@ -13,7 +13,7 @@ import (
type Peer struct { type Peer struct {
id string id string
worker *Worker worker Worker
} }
type WorkerContext struct { type WorkerContext struct {
@@ -74,7 +74,7 @@ func NewPool(ctx context.Context, config *PoolConfig, logger *slog.Logger,
// deadlocks. // deadlocks.
if config.WorkerFactory == nil { if config.WorkerFactory == nil {
config.WorkerFactory = func( config.WorkerFactory = func(
ctx context.Context, id string) (*Worker, error) { ctx context.Context, id string) (Worker, error) {
return NewWorker(ctx, id, config.WorkerConfig) return NewWorker(ctx, id, config.WorkerConfig)
} }
} }
+63 -57
View File
@@ -11,20 +11,26 @@ import (
// Worker // Worker
type receivedMessage struct { type Worker interface {
Start(wctx WorkerContext, wg *sync.WaitGroup)
Stop()
Send(data []byte) error
}
type ReceivedMessage struct {
data []byte data []byte
receivedAt time.Time receivedAt time.Time
} }
type Worker struct { type DefaultWorker struct {
ctx context.Context Ctx context.Context
cancel context.CancelFunc Cancel context.CancelFunc
id string Id string
config *WorkerConfig Config *WorkerConfig
conn atomic.Pointer[transport.Connection] Conn atomic.Pointer[transport.Connection]
heartbeat chan struct{} Heartbeat chan struct{}
} }
func NewWorker( func NewWorker(
@@ -32,7 +38,7 @@ func NewWorker(
id string, id string,
config *WorkerConfig, config *WorkerConfig,
) (*Worker, error) { ) (*DefaultWorker, error) {
if config == nil { if config == nil {
config = GetDefaultWorkerConfig() config = GetDefaultWorkerConfig()
} }
@@ -43,68 +49,68 @@ func NewWorker(
} }
wctx, cancel := context.WithCancel(ctx) wctx, cancel := context.WithCancel(ctx)
w := &Worker{ w := &DefaultWorker{
ctx: wctx, Ctx: wctx,
cancel: cancel, Cancel: cancel,
id: id, Id: id,
config: config, Config: config,
heartbeat: make(chan struct{}), Heartbeat: make(chan struct{}),
} }
return w, nil return w, nil
} }
func (w *Worker) Start( func (w *DefaultWorker) Start(
wctx WorkerContext, wctx WorkerContext,
wg *sync.WaitGroup, wg *sync.WaitGroup,
) { ) {
dial := make(chan struct{}, 1) dial := make(chan struct{}, 1)
newConn := make(chan *transport.Connection, 1) newConn := make(chan *transport.Connection, 1)
messages := make(chan receivedMessage, 256) messages := make(chan ReceivedMessage, 256)
keepalive := make(chan struct{}, 1) keepalive := make(chan struct{}, 1)
var owg sync.WaitGroup var owg sync.WaitGroup
owg.Add(4) owg.Add(4)
go func() { defer owg.Done(); w.runDialer(w.ctx, wctx, dial, newConn) }() go func() { defer owg.Done(); w.RunDialer(w.Ctx, wctx, dial, newConn) }()
go func() { defer owg.Done(); w.runKeepalive(w.ctx, keepalive) }() go func() { defer owg.Done(); w.RunKeepalive(w.Ctx, keepalive) }()
go func() { defer owg.Done(); w.runForwarder(w.ctx, messages, wctx.Inbox, w.config.MaxQueueSize) }() go func() { defer owg.Done(); w.RunForwarder(w.Ctx, messages, wctx.Inbox, w.Config.MaxQueueSize) }()
go func() { defer owg.Done(); w.runSession(w.ctx, wctx, messages, dial, keepalive, newConn) }() go func() { defer owg.Done(); w.RunSession(w.Ctx, wctx, messages, dial, keepalive, newConn) }()
owg.Wait() owg.Wait()
wg.Done() wg.Done()
} }
func (w *Worker) Stop() { func (w *DefaultWorker) Stop() {
w.cancel() w.Cancel()
} }
func (w *Worker) Send(data []byte) error { func (w *DefaultWorker) Send(data []byte) error {
conn := w.conn.Load() conn := w.Conn.Load()
if conn == nil { if conn == nil {
// connection not established by session // connection not established by session
return NewWorkerError(w.id, ErrConnectionUnavailable) return NewWorkerError(w.Id, ErrConnectionUnavailable)
} }
err := conn.Send(data) err := conn.Send(data)
if err != nil { if err != nil {
return NewWorkerError(w.id, err) return NewWorkerError(w.Id, err)
} }
select { select {
case w.heartbeat <- struct{}{}: case w.Heartbeat <- struct{}{}:
case <-w.ctx.Done(): case <-w.Ctx.Done():
} }
return nil return nil
} }
func (w *Worker) runSession( func (w *DefaultWorker) RunSession(
ctx context.Context, ctx context.Context,
wctx WorkerContext, wctx WorkerContext,
messages chan<- receivedMessage, messages chan<- ReceivedMessage,
dial chan<- struct{}, dial chan<- struct{},
keepalive <-chan struct{}, keepalive <-chan struct{},
@@ -135,8 +141,8 @@ func (w *Worker) runSession(
} }
// set up new connection // set up new connection
w.conn.Store(conn) w.Conn.Store(conn)
wctx.Events <- PoolEvent{ID: w.id, Kind: EventConnected} wctx.Events <- PoolEvent{ID: w.Id, Kind: EventConnected}
// set up session // set up session
sessionDone := make(chan struct{}) sessionDone := make(chan struct{})
@@ -150,19 +156,19 @@ func (w *Worker) runSession(
wg.Add(2) wg.Add(2)
go func() { go func() {
defer wg.Done() defer wg.Done()
w.runReader(conn, messages, sessionDone, onStop) w.RunReader(conn, messages, sessionDone, onStop)
}() }()
go func() { go func() {
defer wg.Done() defer wg.Done()
w.runStopMonitor(ctx, conn, keepalive, sessionDone, onStop) w.RunStopMonitor(ctx, conn, keepalive, sessionDone, onStop)
}() }()
// complete session // complete session
wg.Wait() wg.Wait()
// tear down connection // tear down connection
w.conn.Store(nil) w.Conn.Store(nil)
wctx.Events <- PoolEvent{ID: w.id, Kind: EventDisconnected} wctx.Events <- PoolEvent{ID: w.Id, Kind: EventDisconnected}
// exit if worker is shutting down // exit if worker is shutting down
select { select {
@@ -176,9 +182,9 @@ func (w *Worker) runSession(
} }
func (w *Worker) runReader( func (w *DefaultWorker) RunReader(
conn *transport.Connection, conn *transport.Connection,
messages chan<- receivedMessage, messages chan<- ReceivedMessage,
sessionDone <-chan struct{}, sessionDone <-chan struct{},
onStop func(), onStop func(),
) { ) {
@@ -198,14 +204,14 @@ func (w *Worker) runReader(
} }
// send message forward // send message forward
messages <- receivedMessage{ messages <- ReceivedMessage{
data: data, data: data,
receivedAt: time.Now(), receivedAt: time.Now(),
} }
// send heartbeat // send heartbeat
select { select {
case w.heartbeat <- struct{}{}: case w.Heartbeat <- struct{}{}:
case <-sessionDone: case <-sessionDone:
return return
} }
@@ -213,7 +219,7 @@ func (w *Worker) runReader(
} }
} }
func (w *Worker) runStopMonitor( func (w *DefaultWorker) RunStopMonitor(
ctx context.Context, ctx context.Context,
conn *transport.Connection, conn *transport.Connection,
keepalive <-chan struct{}, keepalive <-chan struct{},
@@ -232,9 +238,9 @@ func (w *Worker) runStopMonitor(
} }
} }
func (w *Worker) runForwarder( func (w *DefaultWorker) RunForwarder(
ctx context.Context, ctx context.Context,
messages <-chan receivedMessage, messages <-chan ReceivedMessage,
inbox chan<- InboxMessage, inbox chan<- InboxMessage,
maxQueueSize int, maxQueueSize int,
) { ) {
@@ -242,14 +248,14 @@ func (w *Worker) runForwarder(
for { for {
var out chan<- InboxMessage var out chan<- InboxMessage
var next receivedMessage var next ReceivedMessage
// enable inbox if it is populated // enable inbox if it is populated
if queue.Len() > 0 { if queue.Len() > 0 {
out = inbox out = inbox
// read the first message in the queue // read the first message in the queue
next = queue.Front().Value.(receivedMessage) next = queue.Front().Value.(ReceivedMessage)
} }
select { select {
@@ -265,7 +271,7 @@ func (w *Worker) runForwarder(
queue.PushBack(msg) queue.PushBack(msg)
// send next message to inbox // send next message to inbox
case out <- InboxMessage{ case out <- InboxMessage{
ID: w.id, ID: w.Id,
Data: next.data, Data: next.data,
ReceivedAt: next.receivedAt, ReceivedAt: next.receivedAt,
}: }:
@@ -275,12 +281,12 @@ func (w *Worker) runForwarder(
} }
} }
func (w *Worker) runKeepalive( func (w *DefaultWorker) RunKeepalive(
ctx context.Context, ctx context.Context,
keepalive chan<- struct{}, keepalive chan<- struct{},
) { ) {
// disable keepalive timeout if not configured // disable keepalive timeout if not configured
if w.config.KeepaliveTimeout <= 0 { if w.Config.KeepaliveTimeout <= 0 {
// wait for cancel and exit // wait for cancel and exit
select { select {
case <-ctx.Done(): case <-ctx.Done():
@@ -288,14 +294,14 @@ func (w *Worker) runKeepalive(
return return
} }
timer := time.NewTimer(w.config.KeepaliveTimeout) timer := time.NewTimer(w.Config.KeepaliveTimeout)
defer timer.Stop() defer timer.Stop()
for { for {
select { select {
case <-ctx.Done(): case <-ctx.Done():
return return
case <-w.heartbeat: case <-w.Heartbeat:
// drain the timer channel and reset // drain the timer channel and reset
if !timer.Stop() { if !timer.Stop() {
select { select {
@@ -303,7 +309,7 @@ func (w *Worker) runKeepalive(
default: default:
} }
} }
timer.Reset(w.config.KeepaliveTimeout) timer.Reset(w.Config.KeepaliveTimeout)
// timer completed // timer completed
case <-timer.C: case <-timer.C:
// send keepalive signal, then reset the timer // send keepalive signal, then reset the timer
@@ -311,16 +317,16 @@ func (w *Worker) runKeepalive(
case keepalive <- struct{}{}: case keepalive <- struct{}{}:
default: default:
} }
timer.Reset(w.config.KeepaliveTimeout) timer.Reset(w.Config.KeepaliveTimeout)
} }
} }
} }
func (w *Worker) dial( func (w *DefaultWorker) Dial(
ctx context.Context, ctx context.Context,
wctx WorkerContext, wctx WorkerContext,
) (*transport.Connection, error) { ) (*transport.Connection, error) {
conn, err := transport.NewConnection(w.id, wctx.ConnectionConfig, wctx.Logger) conn, err := transport.NewConnection(w.Id, wctx.ConnectionConfig, wctx.Logger)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@@ -329,7 +335,7 @@ func (w *Worker) dial(
return conn, conn.Connect(ctx) return conn, conn.Connect(ctx)
} }
func (w *Worker) runDialer( func (w *DefaultWorker) RunDialer(
ctx context.Context, ctx context.Context,
wctx WorkerContext, wctx WorkerContext,
@@ -354,7 +360,7 @@ func (w *Worker) runDialer(
}() }()
// dial a new connection // dial a new connection
conn, err := w.dial(ctx, wctx) conn, err := w.Dial(ctx, wctx)
close(done) close(done)
// send error if dial failed and continue // send error if dial failed and continue
+10 -10
View File
@@ -15,7 +15,7 @@ import (
func TestRunDialer(t *testing.T) { func TestRunDialer(t *testing.T) {
t.Run("successful dial delivers connection to newConn", func(t *testing.T) { t.Run("successful dial delivers connection to newConn", func(t *testing.T) {
w := &Worker{id: "wss://test"} w := &DefaultWorker{Id: "wss://test"}
dial := make(chan struct{}, 1) dial := make(chan struct{}, 1)
newConn := make(chan *transport.Connection, 1) newConn := make(chan *transport.Connection, 1)
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
@@ -31,7 +31,7 @@ func TestRunDialer(t *testing.T) {
}, },
} }
go w.runDialer(ctx, wctx, dial, newConn) go w.RunDialer(ctx, wctx, dial, newConn)
dial <- struct{}{} dial <- struct{}{}
honeybeetest.Eventually(t, func() bool { honeybeetest.Eventually(t, func() bool {
@@ -46,7 +46,7 @@ func TestRunDialer(t *testing.T) {
t.Run("concurrent dial signals are drained; only one connection produced.", t.Run("concurrent dial signals are drained; only one connection produced.",
func(t *testing.T) { func(t *testing.T) {
w := &Worker{id: "wss://test"} w := &DefaultWorker{Id: "wss://test"}
dial := make(chan struct{}, 1) dial := make(chan struct{}, 1)
newConn := make(chan *transport.Connection, 1) newConn := make(chan *transport.Connection, 1)
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
@@ -69,7 +69,7 @@ func TestRunDialer(t *testing.T) {
ConnectionConfig: connConfig, ConnectionConfig: connConfig,
} }
go w.runDialer(ctx, wctx, dial, newConn) go w.RunDialer(ctx, wctx, dial, newConn)
dial <- struct{}{} dial <- struct{}{}
// wait for dial to start blocking on gate // wait for dial to start blocking on gate
@@ -107,7 +107,7 @@ func TestRunDialer(t *testing.T) {
}) })
t.Run("dial failure emits error, succeeds on next signal", func(t *testing.T) { t.Run("dial failure emits error, succeeds on next signal", func(t *testing.T) {
w := &Worker{id: "wss://test"} w := &DefaultWorker{Id: "wss://test"}
errors := make(chan error, 1) errors := make(chan error, 1)
dial := make(chan struct{}, 1) dial := make(chan struct{}, 1)
newConn := make(chan *transport.Connection, 1) newConn := make(chan *transport.Connection, 1)
@@ -135,7 +135,7 @@ func TestRunDialer(t *testing.T) {
ConnectionConfig: connConfig, ConnectionConfig: connConfig,
} }
go w.runDialer(ctx, wctx, dial, newConn) go w.RunDialer(ctx, wctx, dial, newConn)
dial <- struct{}{} dial <- struct{}{}
honeybeetest.Eventually(t, func() bool { honeybeetest.Eventually(t, func() bool {
@@ -160,7 +160,7 @@ func TestRunDialer(t *testing.T) {
}) })
t.Run("exits on context cancellation", func(t *testing.T) { t.Run("exits on context cancellation", func(t *testing.T) {
w := &Worker{id: "wss://test"} w := &DefaultWorker{Id: "wss://test"}
dial := make(chan struct{}, 1) dial := make(chan struct{}, 1)
newConn := make(chan *transport.Connection, 1) newConn := make(chan *transport.Connection, 1)
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
@@ -169,7 +169,7 @@ func TestRunDialer(t *testing.T) {
done := make(chan struct{}) done := make(chan struct{})
go func() { go func() {
w.runDialer(ctx, wctx, dial, newConn) w.RunDialer(ctx, wctx, dial, newConn)
close(done) close(done)
}() }()
@@ -186,7 +186,7 @@ func TestRunDialer(t *testing.T) {
}) })
t.Run("context cancelled during in-progress dial exits without delivering connection", func(t *testing.T) { t.Run("context cancelled during in-progress dial exits without delivering connection", func(t *testing.T) {
w := &Worker{id: "wss://test"} w := &DefaultWorker{Id: "wss://test"}
dial := make(chan struct{}, 1) dial := make(chan struct{}, 1)
newConn := make(chan *transport.Connection, 1) newConn := make(chan *transport.Connection, 1)
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
@@ -207,7 +207,7 @@ func TestRunDialer(t *testing.T) {
done := make(chan struct{}) done := make(chan struct{})
go func() { go func() {
w.runDialer(ctx, wctx, dial, newConn) w.RunDialer(ctx, wctx, dial, newConn)
close(done) close(done)
}() }()
+13 -13
View File
@@ -10,15 +10,15 @@ import (
func TestRunForwarder(t *testing.T) { func TestRunForwarder(t *testing.T) {
t.Run("message passes through to inbox", func(t *testing.T) { t.Run("message passes through to inbox", func(t *testing.T) {
messages := make(chan receivedMessage, 1) messages := make(chan ReceivedMessage, 1)
inbox := make(chan InboxMessage, 1) inbox := make(chan InboxMessage, 1)
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
defer cancel() defer cancel()
w := &Worker{id: "wss://test"} w := &DefaultWorker{Id: "wss://test"}
go w.runForwarder(ctx, messages, inbox, 0) go w.RunForwarder(ctx, messages, inbox, 0)
messages <- receivedMessage{data: []byte("hello"), receivedAt: time.Now()} messages <- ReceivedMessage{data: []byte("hello"), receivedAt: time.Now()}
honeybeetest.Eventually(t, func() bool { honeybeetest.Eventually(t, func() bool {
select { select {
@@ -31,7 +31,7 @@ func TestRunForwarder(t *testing.T) {
}) })
t.Run("oldest message dropped when queue is full", func(t *testing.T) { t.Run("oldest message dropped when queue is full", func(t *testing.T) {
messages := make(chan receivedMessage, 1) messages := make(chan ReceivedMessage, 1)
inbox := make(chan InboxMessage, 1) inbox := make(chan InboxMessage, 1)
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
defer cancel() defer cancel()
@@ -47,13 +47,13 @@ func TestRunForwarder(t *testing.T) {
} }
}() }()
w := &Worker{id: "wss://test"} w := &DefaultWorker{Id: "wss://test"}
go w.runForwarder(ctx, messages, gatedInbox, 2) go w.RunForwarder(ctx, messages, gatedInbox, 2)
// send three messages while the gated inbox is blocked // send three messages while the gated inbox is blocked
messages <- receivedMessage{data: []byte("first"), receivedAt: time.Now()} messages <- ReceivedMessage{data: []byte("first"), receivedAt: time.Now()}
messages <- receivedMessage{data: []byte("second"), receivedAt: time.Now()} messages <- ReceivedMessage{data: []byte("second"), receivedAt: time.Now()}
messages <- receivedMessage{data: []byte("third"), receivedAt: time.Now()} messages <- ReceivedMessage{data: []byte("third"), receivedAt: time.Now()}
// allow time for the first message to be dropped // allow time for the first message to be dropped
time.Sleep(20 * time.Millisecond) time.Sleep(20 * time.Millisecond)
@@ -78,15 +78,15 @@ func TestRunForwarder(t *testing.T) {
}) })
t.Run("exits on context cancellation", func(t *testing.T) { t.Run("exits on context cancellation", func(t *testing.T) {
messages := make(chan receivedMessage, 1) messages := make(chan ReceivedMessage, 1)
inbox := make(chan InboxMessage, 1) inbox := make(chan InboxMessage, 1)
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
defer cancel() defer cancel()
w := &Worker{id: "wss://test"} w := &DefaultWorker{Id: "wss://test"}
done := make(chan struct{}) done := make(chan struct{})
go func() { go func() {
w.runForwarder(ctx, messages, inbox, 0) w.RunForwarder(ctx, messages, inbox, 0)
close(done) close(done)
}() }()
+9 -9
View File
@@ -14,16 +14,16 @@ func TestRunKeepalive(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
defer cancel() defer cancel()
w := &Worker{ w := &DefaultWorker{
config: &WorkerConfig{KeepaliveTimeout: 100 * time.Millisecond}, Config: &WorkerConfig{KeepaliveTimeout: 100 * time.Millisecond},
heartbeat: heartbeat, Heartbeat: heartbeat,
} }
go w.runKeepalive(ctx, keepalive) go w.RunKeepalive(ctx, keepalive)
// send heartbeats faster than the timeout // send heartbeats faster than the timeout
for i := 0; i < 5; i++ { for i := 0; i < 5; i++ {
time.Sleep(30 * time.Millisecond) time.Sleep(30 * time.Millisecond)
w.heartbeat <- struct{}{} w.Heartbeat <- struct{}{}
} }
// because the timer is being reset, keepalive signal should not be sent // because the timer is being reset, keepalive signal should not be sent
@@ -42,8 +42,8 @@ func TestRunKeepalive(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
defer cancel() defer cancel()
w := &Worker{config: &WorkerConfig{KeepaliveTimeout: 20 * time.Millisecond}} w := &DefaultWorker{Config: &WorkerConfig{KeepaliveTimeout: 20 * time.Millisecond}}
go w.runKeepalive(ctx, keepalive) go w.RunKeepalive(ctx, keepalive)
// send no heartbeats, wait for timeout and keepalive signal // send no heartbeats, wait for timeout and keepalive signal
honeybeetest.Eventually(t, func() bool { honeybeetest.Eventually(t, func() bool {
@@ -60,10 +60,10 @@ func TestRunKeepalive(t *testing.T) {
keepalive := make(chan struct{}, 1) keepalive := make(chan struct{}, 1)
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
w := &Worker{config: &WorkerConfig{KeepaliveTimeout: 20 * time.Second}} w := &DefaultWorker{Config: &WorkerConfig{KeepaliveTimeout: 20 * time.Second}}
done := make(chan struct{}) done := make(chan struct{})
go func() { go func() {
w.runKeepalive(ctx, keepalive) w.RunKeepalive(ctx, keepalive)
close(done) close(done)
}() }()
+20 -20
View File
@@ -19,14 +19,14 @@ func TestWorkerSend(t *testing.T) {
heartbeat := make(chan struct{}) heartbeat := make(chan struct{})
heartbeatCount := atomic.Int32{} heartbeatCount := atomic.Int32{}
w := &Worker{ w := &DefaultWorker{
ctx: ctx, Ctx: ctx,
cancel: cancel, Cancel: cancel,
id: "wss://test", Id: "wss://test",
heartbeat: heartbeat, Heartbeat: heartbeat,
} }
w.conn.Store(conn) w.Conn.Store(conn)
defer w.cancel() defer w.Cancel()
go func() { go func() {
for range heartbeat { for range heartbeat {
@@ -61,14 +61,14 @@ func TestWorkerSend(t *testing.T) {
heartbeat := make(chan struct{}) heartbeat := make(chan struct{})
heartbeatCount := atomic.Int32{} heartbeatCount := atomic.Int32{}
w := &Worker{ w := &DefaultWorker{
ctx: ctx, Ctx: ctx,
cancel: cancel, Cancel: cancel,
id: "wss://test", Id: "wss://test",
heartbeat: heartbeat, Heartbeat: heartbeat,
} }
w.conn.Store(conn) w.Conn.Store(conn)
defer w.cancel() defer w.Cancel()
go func() { go func() {
for range heartbeat { for range heartbeat {
@@ -92,13 +92,13 @@ func TestWorkerSend(t *testing.T) {
heartbeat := make(chan struct{}) heartbeat := make(chan struct{})
w := &Worker{ w := &DefaultWorker{
ctx: ctx, Ctx: ctx,
cancel: cancel, Cancel: cancel,
id: "wss://test", Id: "wss://test",
heartbeat: heartbeat, Heartbeat: heartbeat,
} }
defer w.cancel() defer w.Cancel()
go func() { go func() {
for range heartbeat { for range heartbeat {
+32 -32
View File
@@ -18,24 +18,24 @@ func TestRunReader(t *testing.T) {
conn, _, incomingData, _ := setupWorkerTestConnection(t) conn, _, incomingData, _ := setupWorkerTestConnection(t)
defer conn.Close() defer conn.Close()
messages := make(chan receivedMessage, 1) messages := make(chan ReceivedMessage, 1)
heartbeat := make(chan struct{}) heartbeat := make(chan struct{})
sessionDone := make(chan struct{}) sessionDone := make(chan struct{})
onStop := func() {} onStop := func() {}
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
defer cancel() defer cancel()
w := &Worker{ w := &DefaultWorker{
ctx: ctx, Ctx: ctx,
cancel: cancel, Cancel: cancel,
id: "wss://test", Id: "wss://test",
heartbeat: heartbeat, Heartbeat: heartbeat,
} }
go func() { go func() {
for range heartbeat { for range heartbeat {
} }
}() }()
go w.runReader(conn, messages, sessionDone, onStop) go w.RunReader(conn, messages, sessionDone, onStop)
before := time.Now() before := time.Now()
incomingData <- honeybeetest.MockIncomingData{ incomingData <- honeybeetest.MockIncomingData{
@@ -57,18 +57,18 @@ func TestRunReader(t *testing.T) {
conn, _, incomingData, _ := setupWorkerTestConnection(t) conn, _, incomingData, _ := setupWorkerTestConnection(t)
defer conn.Close() defer conn.Close()
messages := make(chan receivedMessage, 10) messages := make(chan ReceivedMessage, 10)
heartbeat := make(chan struct{}) heartbeat := make(chan struct{})
sessionDone := make(chan struct{}) sessionDone := make(chan struct{})
onStop := func() {} onStop := func() {}
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
defer cancel() defer cancel()
w := &Worker{ w := &DefaultWorker{
ctx: ctx, Ctx: ctx,
cancel: cancel, Cancel: cancel,
id: "wss://test", Id: "wss://test",
heartbeat: heartbeat, Heartbeat: heartbeat,
} }
received := atomic.Int32{} received := atomic.Int32{}
@@ -81,7 +81,7 @@ func TestRunReader(t *testing.T) {
for range messages { for range messages {
} }
}() }()
go w.runReader(conn, messages, sessionDone, onStop) go w.RunReader(conn, messages, sessionDone, onStop)
const count = 3 const count = 3
for i := 0; i < count; i++ { for i := 0; i < count; i++ {
@@ -99,17 +99,17 @@ func TestRunReader(t *testing.T) {
t.Run("incoming channel close calls conn.Close and onStop", func(t *testing.T) { t.Run("incoming channel close calls conn.Close and onStop", func(t *testing.T) {
conn, _, incomingData, _ := setupWorkerTestConnection(t) conn, _, incomingData, _ := setupWorkerTestConnection(t)
messages := make(chan receivedMessage, 1) messages := make(chan ReceivedMessage, 1)
heartbeat := make(chan struct{}) heartbeat := make(chan struct{})
sessionDone := make(chan struct{}) sessionDone := make(chan struct{})
onStopCalled := atomic.Bool{} onStopCalled := atomic.Bool{}
onStop := func() { onStopCalled.Store(true) } onStop := func() { onStopCalled.Store(true) }
ctx := context.Background() ctx := context.Background()
w := &Worker{ w := &DefaultWorker{
ctx: ctx, Ctx: ctx,
id: "wss://test", Id: "wss://test",
heartbeat: heartbeat, Heartbeat: heartbeat,
} }
go func() { go func() {
for range heartbeat { for range heartbeat {
@@ -119,7 +119,7 @@ func TestRunReader(t *testing.T) {
for range messages { for range messages {
} }
}() }()
go w.runReader(conn, messages, sessionDone, onStop) go w.RunReader(conn, messages, sessionDone, onStop)
// induce connection closure via reader // induce connection closure via reader
incomingData <- honeybeetest.MockIncomingData{Err: io.EOF} incomingData <- honeybeetest.MockIncomingData{Err: io.EOF}
@@ -139,19 +139,19 @@ func TestRunReader(t *testing.T) {
t.Run("sessionDone close calls conn.Close and onStop", func(t *testing.T) { t.Run("sessionDone close calls conn.Close and onStop", func(t *testing.T) {
conn, _, _, _ := setupWorkerTestConnection(t) conn, _, _, _ := setupWorkerTestConnection(t)
messages := make(chan receivedMessage, 1) messages := make(chan ReceivedMessage, 1)
heartbeat := make(chan struct{}) heartbeat := make(chan struct{})
sessionDone := make(chan struct{}) sessionDone := make(chan struct{})
onStopCalled := atomic.Bool{} onStopCalled := atomic.Bool{}
onStop := func() { onStopCalled.Store(true) } onStop := func() { onStopCalled.Store(true) }
ctx := context.Background() ctx := context.Background()
w := &Worker{ w := &DefaultWorker{
ctx: ctx, Ctx: ctx,
id: "wss://test", Id: "wss://test",
heartbeat: heartbeat, Heartbeat: heartbeat,
} }
go w.runReader(conn, messages, sessionDone, onStop) go w.RunReader(conn, messages, sessionDone, onStop)
close(sessionDone) close(sessionDone)
@@ -176,8 +176,8 @@ func TestRunStopMonitor(t *testing.T) {
onStopCalled := atomic.Bool{} onStopCalled := atomic.Bool{}
onStop := func() { onStopCalled.Store(true) } onStop := func() { onStopCalled.Store(true) }
w := &Worker{id: "wss://test"} w := &DefaultWorker{Id: "wss://test"}
go w.runStopMonitor(ctx, conn, keepalive, sessionDone, onStop) go w.RunStopMonitor(ctx, conn, keepalive, sessionDone, onStop)
keepalive <- struct{}{} keepalive <- struct{}{}
@@ -199,8 +199,8 @@ func TestRunStopMonitor(t *testing.T) {
onStopCalled := atomic.Bool{} onStopCalled := atomic.Bool{}
onStop := func() { onStopCalled.Store(true) } onStop := func() { onStopCalled.Store(true) }
w := &Worker{id: "wss://test"} w := &DefaultWorker{Id: "wss://test"}
go w.runStopMonitor(ctx, conn, keepalive, sessionDone, onStop) go w.RunStopMonitor(ctx, conn, keepalive, sessionDone, onStop)
cancel() cancel()
@@ -223,8 +223,8 @@ func TestRunStopMonitor(t *testing.T) {
onStopCalled := atomic.Bool{} onStopCalled := atomic.Bool{}
onStop := func() { onStopCalled.Store(true) } onStop := func() { onStopCalled.Store(true) }
w := &Worker{id: "wss://test"} w := &DefaultWorker{Id: "wss://test"}
go w.runStopMonitor(ctx, conn, keepalive, sessionDone, onStop) go w.RunStopMonitor(ctx, conn, keepalive, sessionDone, onStop)
close(sessionDone) close(sessionDone)
+53 -53
View File
@@ -22,7 +22,7 @@ func drainEvent(t *testing.T, events <-chan PoolEvent, kind PoolEventKind) {
func TestRunSessionDial(t *testing.T) { func TestRunSessionDial(t *testing.T) {
setup := func(t *testing.T) ( setup := func(t *testing.T) (
w *Worker, w *DefaultWorker,
ctx context.Context, ctx context.Context,
cancel context.CancelFunc, cancel context.CancelFunc,
dial chan struct{}, dial chan struct{},
@@ -31,12 +31,12 @@ func TestRunSessionDial(t *testing.T) {
) { ) {
t.Helper() t.Helper()
ctx, cancel = context.WithCancel(context.Background()) ctx, cancel = context.WithCancel(context.Background())
w = &Worker{ w = &DefaultWorker{
ctx: ctx, Ctx: ctx,
cancel: cancel, Cancel: cancel,
id: "wss://test", Id: "wss://test",
config: GetDefaultWorkerConfig(), Config: GetDefaultWorkerConfig(),
heartbeat: make(chan struct{}), Heartbeat: make(chan struct{}),
} }
dial = make(chan struct{}, 1) dial = make(chan struct{}, 1)
keepalive = make(chan struct{}, 1) keepalive = make(chan struct{}, 1)
@@ -60,10 +60,10 @@ func TestRunSessionDial(t *testing.T) {
w, ctx, cancel, dial, keepalive, newConn := setup(t) w, ctx, cancel, dial, keepalive, newConn := setup(t)
defer cancel() defer cancel()
messages := make(chan receivedMessage, 1) messages := make(chan ReceivedMessage, 1)
wctx := WorkerContext{Events: make(chan PoolEvent, 10)} wctx := WorkerContext{Events: make(chan PoolEvent, 10)}
go w.runSession(ctx, wctx, messages, dial, keepalive, newConn) go w.RunSession(ctx, wctx, messages, dial, keepalive, newConn)
expectDial(t, dial) expectDial(t, dial)
}) })
@@ -72,10 +72,10 @@ func TestRunSessionDial(t *testing.T) {
w, ctx, cancel, dial, keepalive, newConn := setup(t) w, ctx, cancel, dial, keepalive, newConn := setup(t)
defer cancel() defer cancel()
messages := make(chan receivedMessage, 1) messages := make(chan ReceivedMessage, 1)
wctx := WorkerContext{Events: make(chan PoolEvent, 10)} wctx := WorkerContext{Events: make(chan PoolEvent, 10)}
go w.runSession(ctx, wctx, messages, dial, keepalive, newConn) go w.RunSession(ctx, wctx, messages, dial, keepalive, newConn)
// drain initial dial // drain initial dial
expectDial(t, dial) expectDial(t, dial)
@@ -88,10 +88,10 @@ func TestRunSessionDial(t *testing.T) {
w, ctx, cancel, dial, keepalive, newConn := setup(t) w, ctx, cancel, dial, keepalive, newConn := setup(t)
defer cancel() defer cancel()
messages := make(chan receivedMessage, 1) messages := make(chan ReceivedMessage, 1)
wctx := WorkerContext{Events: make(chan PoolEvent, 10)} wctx := WorkerContext{Events: make(chan PoolEvent, 10)}
go w.runSession(ctx, wctx, messages, dial, keepalive, newConn) go w.RunSession(ctx, wctx, messages, dial, keepalive, newConn)
// drain initial dial // drain initial dial
expectDial(t, dial) expectDial(t, dial)
@@ -105,27 +105,27 @@ func TestRunSessionDial(t *testing.T) {
func TestRunSessionConnect(t *testing.T) { func TestRunSessionConnect(t *testing.T) {
setup := func(t *testing.T) ( setup := func(t *testing.T) (
w *Worker, w *DefaultWorker,
ctx context.Context, ctx context.Context,
cancel context.CancelFunc, cancel context.CancelFunc,
dial chan struct{}, dial chan struct{},
keepalive chan struct{}, keepalive chan struct{},
newConn chan *transport.Connection, newConn chan *transport.Connection,
messages chan receivedMessage, messages chan ReceivedMessage,
) { ) {
t.Helper() t.Helper()
ctx, cancel = context.WithCancel(context.Background()) ctx, cancel = context.WithCancel(context.Background())
w = &Worker{ w = &DefaultWorker{
ctx: ctx, Ctx: ctx,
cancel: cancel, Cancel: cancel,
id: "wss://test", Id: "wss://test",
config: GetDefaultWorkerConfig(), Config: GetDefaultWorkerConfig(),
heartbeat: make(chan struct{}), Heartbeat: make(chan struct{}),
} }
dial = make(chan struct{}, 1) dial = make(chan struct{}, 1)
keepalive = make(chan struct{}, 1) keepalive = make(chan struct{}, 1)
newConn = make(chan *transport.Connection, 1) newConn = make(chan *transport.Connection, 1)
messages = make(chan receivedMessage, 256) messages = make(chan ReceivedMessage, 256)
return return
} }
@@ -135,12 +135,12 @@ func TestRunSessionConnect(t *testing.T) {
defer cancel() defer cancel()
conn, _, _, _ := setupWorkerTestConnection(t) conn, _, _, _ := setupWorkerTestConnection(t)
go w.runSession(ctx, wctx, messages, dial, keepalive, newConn) go w.RunSession(ctx, wctx, messages, dial, keepalive, newConn)
newConn <- conn newConn <- conn
honeybeetest.Eventually(t, func() bool { honeybeetest.Eventually(t, func() bool {
return w.conn.Load() != nil return w.Conn.Load() != nil
}, "expected w.conn to be set") }, "expected w.conn to be set")
}) })
@@ -151,14 +151,14 @@ func TestRunSessionConnect(t *testing.T) {
defer cancel() defer cancel()
conn, _, _, _ := setupWorkerTestConnection(t) conn, _, _, _ := setupWorkerTestConnection(t)
go w.runSession(ctx, wctx, messages, dial, keepalive, newConn) go w.RunSession(ctx, wctx, messages, dial, keepalive, newConn)
newConn <- conn newConn <- conn
honeybeetest.Eventually(t, func() bool { honeybeetest.Eventually(t, func() bool {
select { select {
case event := <-events: case event := <-events:
return event.ID == w.id && event.Kind == EventConnected return event.ID == w.Id && event.Kind == EventConnected
default: default:
return false return false
} }
@@ -168,29 +168,29 @@ func TestRunSessionConnect(t *testing.T) {
func TestRunSessionDisconnect(t *testing.T) { func TestRunSessionDisconnect(t *testing.T) {
setup := func(t *testing.T) ( setup := func(t *testing.T) (
w *Worker, w *DefaultWorker,
ctx context.Context, ctx context.Context,
cancel context.CancelFunc, cancel context.CancelFunc,
dial chan struct{}, dial chan struct{},
keepalive chan struct{}, keepalive chan struct{},
newConn chan *transport.Connection, newConn chan *transport.Connection,
messages chan receivedMessage, messages chan ReceivedMessage,
conn *transport.Connection, conn *transport.Connection,
incomingData chan honeybeetest.MockIncomingData, incomingData chan honeybeetest.MockIncomingData,
) { ) {
t.Helper() t.Helper()
ctx, cancel = context.WithCancel(context.Background()) ctx, cancel = context.WithCancel(context.Background())
w = &Worker{ w = &DefaultWorker{
ctx: ctx, Ctx: ctx,
cancel: cancel, Cancel: cancel,
id: "wss://test", Id: "wss://test",
config: GetDefaultWorkerConfig(), Config: GetDefaultWorkerConfig(),
heartbeat: make(chan struct{}), Heartbeat: make(chan struct{}),
} }
dial = make(chan struct{}, 1) dial = make(chan struct{}, 1)
keepalive = make(chan struct{}, 1) keepalive = make(chan struct{}, 1)
newConn = make(chan *transport.Connection, 1) newConn = make(chan *transport.Connection, 1)
messages = make(chan receivedMessage, 256) messages = make(chan ReceivedMessage, 256)
conn, _, incomingData, _ = setupWorkerTestConnection(t) conn, _, incomingData, _ = setupWorkerTestConnection(t)
return return
} }
@@ -201,7 +201,7 @@ func TestRunSessionDisconnect(t *testing.T) {
wctx := WorkerContext{Events: events} wctx := WorkerContext{Events: events}
defer cancel() defer cancel()
go w.runSession(ctx, wctx, messages, dial, keepalive, newConn) go w.RunSession(ctx, wctx, messages, dial, keepalive, newConn)
newConn <- conn newConn <- conn
drainEvent(t, events, EventConnected) drainEvent(t, events, EventConnected)
@@ -216,7 +216,7 @@ func TestRunSessionDisconnect(t *testing.T) {
wctx := WorkerContext{Events: events} wctx := WorkerContext{Events: events}
defer cancel() defer cancel()
go w.runSession(ctx, wctx, messages, dial, keepalive, newConn) go w.RunSession(ctx, wctx, messages, dial, keepalive, newConn)
newConn <- conn newConn <- conn
drainEvent(t, events, EventConnected) drainEvent(t, events, EventConnected)
@@ -224,7 +224,7 @@ func TestRunSessionDisconnect(t *testing.T) {
drainEvent(t, events, EventDisconnected) drainEvent(t, events, EventDisconnected)
honeybeetest.Eventually(t, func() bool { honeybeetest.Eventually(t, func() bool {
return w.conn.Load() == nil return w.Conn.Load() == nil
}, "expected w.conn to be cleared") }, "expected w.conn to be cleared")
}) })
@@ -234,7 +234,7 @@ func TestRunSessionDisconnect(t *testing.T) {
wctx := WorkerContext{Events: events} wctx := WorkerContext{Events: events}
defer cancel() defer cancel()
go w.runSession(ctx, wctx, messages, dial, keepalive, newConn) go w.RunSession(ctx, wctx, messages, dial, keepalive, newConn)
newConn <- conn newConn <- conn
drainEvent(t, events, EventConnected) drainEvent(t, events, EventConnected)
@@ -260,7 +260,7 @@ func TestRunSessionDisconnect(t *testing.T) {
wctx := WorkerContext{Events: events} wctx := WorkerContext{Events: events}
defer cancel() defer cancel()
go w.runSession(ctx, wctx, messages, dial, keepalive, newConn) go w.RunSession(ctx, wctx, messages, dial, keepalive, newConn)
newConn <- conn newConn <- conn
drainEvent(t, events, EventConnected) drainEvent(t, events, EventConnected)
@@ -276,27 +276,27 @@ func TestRunSessionDisconnect(t *testing.T) {
func TestRunSessionCancellation(t *testing.T) { func TestRunSessionCancellation(t *testing.T) {
setup := func(t *testing.T) ( setup := func(t *testing.T) (
w *Worker, w *DefaultWorker,
ctx context.Context, ctx context.Context,
cancel context.CancelFunc, cancel context.CancelFunc,
dial chan struct{}, dial chan struct{},
keepalive chan struct{}, keepalive chan struct{},
newConn chan *transport.Connection, newConn chan *transport.Connection,
messages chan receivedMessage, messages chan ReceivedMessage,
) { ) {
t.Helper() t.Helper()
ctx, cancel = context.WithCancel(context.Background()) ctx, cancel = context.WithCancel(context.Background())
w = &Worker{ w = &DefaultWorker{
ctx: ctx, Ctx: ctx,
cancel: cancel, Cancel: cancel,
id: "wss://test", Id: "wss://test",
config: GetDefaultWorkerConfig(), Config: GetDefaultWorkerConfig(),
heartbeat: make(chan struct{}), Heartbeat: make(chan struct{}),
} }
dial = make(chan struct{}, 1) dial = make(chan struct{}, 1)
keepalive = make(chan struct{}, 1) keepalive = make(chan struct{}, 1)
newConn = make(chan *transport.Connection, 1) newConn = make(chan *transport.Connection, 1)
messages = make(chan receivedMessage, 256) messages = make(chan ReceivedMessage, 256)
return return
} }
@@ -308,7 +308,7 @@ func TestRunSessionCancellation(t *testing.T) {
done := make(chan struct{}) done := make(chan struct{})
go func() { go func() {
defer close(done) defer close(done)
w.runSession(ctx, wctx, messages, dial, keepalive, newConn) w.RunSession(ctx, wctx, messages, dial, keepalive, newConn)
}() }()
cancel() cancel()
@@ -342,7 +342,7 @@ func TestRunSessionCancellation(t *testing.T) {
done := make(chan struct{}) done := make(chan struct{})
go func() { go func() {
defer close(done) defer close(done)
w.runSession(ctx, wctx, messages, dial, keepalive, newConn) w.RunSession(ctx, wctx, messages, dial, keepalive, newConn)
}() }()
newConn <- conn newConn <- conn
@@ -373,7 +373,7 @@ func TestRunSessionCancellation(t *testing.T) {
done := make(chan struct{}) done := make(chan struct{})
go func() { go func() {
defer close(done) defer close(done)
w.runSession(ctx, wctx, messages, dial, keepalive, newConn) w.RunSession(ctx, wctx, messages, dial, keepalive, newConn)
}() }()
newConn <- conn newConn <- conn
@@ -385,7 +385,7 @@ func TestRunSessionCancellation(t *testing.T) {
drainEvent(t, events, EventDisconnected) drainEvent(t, events, EventDisconnected)
honeybeetest.Eventually(t, func() bool { honeybeetest.Eventually(t, func() bool {
return w.conn.Load() == nil return w.Conn.Load() == nil
}, "expected w.conn to clear") }, "expected w.conn to clear")
}) })
} }
+9 -9
View File
@@ -31,14 +31,14 @@ func makeWorkerContext(t *testing.T) (
return return
} }
func makeWorker(t *testing.T, ctx context.Context, cancel context.CancelFunc) *Worker { func makeWorker(t *testing.T, ctx context.Context, cancel context.CancelFunc) *DefaultWorker {
t.Helper() t.Helper()
return &Worker{ return &DefaultWorker{
ctx: ctx, Ctx: ctx,
cancel: cancel, Cancel: cancel,
id: "wss://test", Id: "wss://test",
config: GetDefaultWorkerConfig(), Config: GetDefaultWorkerConfig(),
heartbeat: make(chan struct{}), Heartbeat: make(chan struct{}),
} }
} }
@@ -67,7 +67,7 @@ func TestWorkerStart(t *testing.T) {
honeybeetest.Eventually(t, func() bool { honeybeetest.Eventually(t, func() bool {
select { select {
case e := <-events: case e := <-events:
return e.ID == w.id && e.Kind == EventConnected return e.ID == w.Id && e.Kind == EventConnected
default: default:
return false return false
} }
@@ -154,7 +154,7 @@ func TestWorkerStart(t *testing.T) {
honeybeetest.Eventually(t, func() bool { honeybeetest.Eventually(t, func() bool {
select { select {
case msg := <-inbox: case msg := <-inbox:
return msg.ID == w.id && string(msg.Data) == "hello" return msg.ID == w.Id && string(msg.Data) == "hello"
default: default:
return false return false
} }