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

14
transport/connection.go
View File

@@ -1,6 +1,7 @@
package transport
import (
"context"
"errors"
"fmt"
"log/slog"
@@ -64,13 +65,13 @@ func NewConnection(urlStr string, config *ConnectionConfig, logger *slog.Logger)
return nil, err
}
parsedURL, err := ParseURL(urlStr)
url, err := ParseURL(urlStr)
if err != nil {
return nil, err
}
conn := &Connection{
url: parsedURL,
url: url,
dialer: NewDialer(),
socket: nil,
config: config,
@@ -85,7 +86,9 @@ func NewConnection(urlStr string, config *ConnectionConfig, logger *slog.Logger)
return conn, nil
}
func NewConnectionFromSocket(socket types.Socket, config *ConnectionConfig, logger *slog.Logger) (*Connection, error) {
func NewConnectionFromSocket(
socket types.Socket, config *ConnectionConfig, logger *slog.Logger,
) (*Connection, error) {
if socket == nil {
return nil, NewConnectionError("socket cannot be nil")
}
@@ -121,7 +124,7 @@ func NewConnectionFromSocket(socket types.Socket, config *ConnectionConfig, logg
return conn, nil
}
func (c *Connection) Connect() error {
func (c *Connection) Connect(ctx context.Context) error {
c.mu.Lock()
defer c.mu.Unlock()
@@ -140,7 +143,8 @@ func (c *Connection) Connect() error {
c.state = StateConnecting
retryMgr := NewRetryManager(c.config.Retry)
socket, _, err := AcquireSocket(retryMgr, c.dialer, c.url.String(), c.logger)
socket, _, err := AcquireSocket(
ctx, retryMgr, c.dialer, c.url.String(), c.logger)
if err != nil {
c.state = StateDisconnected

25
transport/connection_test.go
View File

@@ -2,6 +2,7 @@ package transport
import (
"bytes"
"context"
"fmt"
"git.wisehodl.dev/jay/go-honeybee/honeybeetest"
"git.wisehodl.dev/jay/go-honeybee/types"
@@ -239,7 +240,7 @@ func TestConnect(t *testing.T) {
conn.socket = honeybeetest.NewMockSocket()
err = conn.Connect()
err = conn.Connect(context.Background())
assert.Error(t, err)
assert.ErrorContains(t, err, "already has socket")
assert.Equal(t, StateDisconnected, conn.State())
@@ -251,7 +252,7 @@ func TestConnect(t *testing.T) {
conn.Close()
err = conn.Connect()
err = conn.Connect(context.Background())
assert.Error(t, err)
assert.ErrorContains(t, err, "connection is closed")
assert.Equal(t, StateClosed, conn.State())
@@ -270,13 +271,13 @@ func TestConnect(t *testing.T) {
}
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
},
}
conn.dialer = mockDialer
err = conn.Connect()
err = conn.Connect(context.Background())
assert.NoError(t, err)
assert.Equal(t, StateConnected, conn.State())
@@ -309,7 +310,7 @@ func TestConnect(t *testing.T) {
attemptCount := 0
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) {
attemptCount++
if attemptCount < 3 {
return nil, nil, fmt.Errorf("dial failed")
@@ -319,7 +320,7 @@ func TestConnect(t *testing.T) {
}
conn.dialer = mockDialer
err = conn.Connect()
err = conn.Connect(context.Background())
assert.NoError(t, err)
assert.Equal(t, 3, attemptCount)
assert.Equal(t, StateConnected, conn.State())
@@ -340,13 +341,13 @@ func TestConnect(t *testing.T) {
assert.NoError(t, err)
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 nil, nil, fmt.Errorf("dial failed")
},
}
conn.dialer = mockDialer
err = conn.Connect()
err = conn.Connect(context.Background())
assert.Error(t, err)
assert.ErrorContains(t, err, "dial failed")
assert.Equal(t, StateDisconnected, conn.State())
@@ -359,14 +360,14 @@ func TestConnect(t *testing.T) {
stateDuringDial := StateDisconnected
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) {
stateDuringDial = conn.state
return honeybeetest.NewMockSocket(), nil, nil
},
}
conn.dialer = mockDialer
conn.Connect()
conn.Connect(context.Background())
assert.Equal(t, StateConnecting, stateDuringDial)
assert.Equal(t, StateConnected, conn.State())
@@ -390,13 +391,13 @@ func TestConnect(t *testing.T) {
}
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
},
}
conn.dialer = mockDialer
conn.Connect()
conn.Connect(context.Background())
assert.True(t, handlerSet, "close handler should be set on socket")

17
transport/logging_test.go
View File

@@ -1,6 +1,7 @@
package transport
import (
"context"
"fmt"
"io"
"log/slog"
@@ -146,13 +147,13 @@ func TestConnectLogging(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
},
}
conn.dialer = mockDialer
err = conn.Connect()
err = conn.Connect(context.Background())
assert.NoError(t, err)
defer conn.Close()
@@ -186,13 +187,13 @@ func TestConnectLogging(t *testing.T) {
dialErr := fmt.Errorf("dial error")
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 nil, nil, dialErr
},
}
conn.dialer = mockDialer
err = conn.Connect()
err = conn.Connect(context.Background())
assert.Error(t, err)
records := mockHandler.GetRecords()
@@ -230,7 +231,7 @@ func TestConnectLogging(t *testing.T) {
attemptCount := 0
dialErr := fmt.Errorf("dial error")
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) {
attemptCount++
if attemptCount < 3 {
return nil, nil, dialErr
@@ -240,7 +241,7 @@ func TestConnectLogging(t *testing.T) {
}
conn.dialer = mockDialer
err = conn.Connect()
err = conn.Connect(context.Background())
assert.NoError(t, err)
defer conn.Close()
@@ -468,13 +469,13 @@ func TestLoggingDisabled(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
},
}
conn.dialer = mockDialer
err = conn.Connect()
err = conn.Connect(context.Background())
assert.NoError(t, err)
conn.Close()

22
transport/socket.go
View File

@@ -1,6 +1,7 @@
package transport
import (
"context"
"log/slog"
"net/http"
"time"
@@ -28,19 +29,22 @@ func NewGorillaDialer() *GorillaDialer {
}
// Returns the Socket interface
func (d *GorillaDialer) Dial(
urlStr string, requestHeader http.Header,
func (d *GorillaDialer) DialContext(
ctx context.Context,
url string,
header http.Header,
) (
types.Socket, *http.Response, error,
) {
conn, resp, err := d.Dialer.Dial(urlStr, requestHeader)
conn, resp, err := d.Dialer.DialContext(ctx, url, header)
return conn, resp, err
}
func AcquireSocket(
ctx context.Context,
retryMgr *RetryManager,
dialer types.Dialer,
urlStr string,
url string,
logger *slog.Logger,
) (types.Socket, *http.Response, error) {
if retryMgr == nil {
@@ -49,7 +53,7 @@ func AcquireSocket(
if dialer == nil {
return nil, nil, NewConnectionError("dialer cannot be nil")
}
if urlStr == "" {
if url == "" {
return nil, nil, NewConnectionError("URL cannot be empty")
}
@@ -58,7 +62,7 @@ func AcquireSocket(
logger.Info("dialing", "attempt", retryMgr.RetryCount()+1)
}
socket, resp, err := dialer.Dial(urlStr, nil)
socket, resp, err := dialer.DialContext(ctx, url, nil)
if err == nil {
if logger != nil {
logger.Info("dial successful", "attempt", retryMgr.RetryCount()+1)
@@ -84,7 +88,11 @@ func AcquireSocket(
"next_delay", delay)
}
time.Sleep(delay)
select {
case <-time.After(delay):
case <-ctx.Done():
return nil, nil, ctx.Err()
}
retryMgr.RecordRetry()
}
}

13
transport/socket_test.go
View File

@@ -1,6 +1,7 @@
package transport
import (
"context"
"errors"
"git.wisehodl.dev/jay/go-honeybee/honeybeetest"
"git.wisehodl.dev/jay/go-honeybee/types"
@@ -63,7 +64,8 @@ func TestAcquireSocket(t *testing.T) {
t.Run(tc.name, func(t *testing.T) {
attemptIndex := 0
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) {
err := tc.mockRuns[attemptIndex]
attemptIndex++
if err != nil {
@@ -80,7 +82,8 @@ func TestAcquireSocket(t *testing.T) {
JitterFactor: 0.0,
})
socket, _, err := AcquireSocket(retryMgr, mockDialer, "ws://test", nil)
socket, _, err := AcquireSocket(
context.Background(), retryMgr, mockDialer, "ws://test", nil)
assert.Equal(t, tc.wantRetryCount, retryMgr.RetryCount())
if tc.wantErr {
@@ -96,7 +99,8 @@ func TestAcquireSocket(t *testing.T) {
func TestAcquireSocketGuards(t *testing.T) {
validDialer := &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 honeybeetest.NewMockSocket(), nil, nil
},
}
@@ -134,7 +138,8 @@ func TestAcquireSocketGuards(t *testing.T) {
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
socket, resp, err := AcquireSocket(tc.retryMgr, tc.dialer, tc.url, nil)
socket, resp, err := AcquireSocket(
context.Background(), tc.retryMgr, tc.dialer, tc.url, nil)
assert.Error(t, err)
assert.ErrorContains(t, err, tc.wantErr)