diff --git a/conn/connpool/onetimepool.go b/conn/connpool/onetimepool.go index 73282f1..8f1d595 100644 --- a/conn/connpool/onetimepool.go +++ b/conn/connpool/onetimepool.go @@ -2,8 +2,17 @@ package connpool import ( "context" + "fmt" "net" "sync" + "time" +) + +// Bounds of the wait between failed dials in OneTimePool.Run. +const ( + minDialBackoff = 50 * time.Millisecond + maxDialBackoff = time.Second + backoffFactor = 2 ) type Dialer interface { @@ -71,35 +80,49 @@ func (p *OneTimePool) Acquire(c context.Context) (net.Conn, error) { } } +// Run dials connections and keeps the pool filled until the pool is closed or +// its context is done, then returns an error wrapping the context's error +// (match it with errors.Is). Failed dials are retried with an exponentially +// growing wait, capped at maxDialBackoff and reset by the next successful +// dial, so an unreachable target does not consume a CPU core. Dial errors are +// not reported; callers observe them as Acquire blocking until its own +// context is done. func (p *OneTimePool) Run() error { + backoff := minDialBackoff for { + if p.ctx.Err() != nil { + return p.stopped() + } + + conn := p.dial() + if conn == nil { + if !p.sleep(backoff) { + return p.stopped() + } + + backoff = min(backoff*backoffFactor, maxDialBackoff) + + continue + } + + backoff = minDialBackoff + + p.InFlightConns.Add(conn) + // Close cancels the context before it closes the in-flight + // connections. A connection added after that sweep would never be + // closed, so a done context here means this goroutine must close it. + if p.ctx.Err() != nil { + p.discard(conn) + + return p.stopped() + } + select { case <-p.ctx.Done(): - return p.ctx.Err() - default: - var ( - conn net.Conn - err error - ) - p.mx.RLock() - hasDialer := p.Dialer != nil - p.mx.RUnlock() - - if hasDialer { - p.mx.RLock() - conn, err = p.Dialer.Dial(p.ctx, "tcp", p.address) - p.mx.RUnlock() - } else { - conn, err = net.Dial("tcp", p.address) - } - if err == nil { - p.InFlightConns.Add(conn) - select { - case <-p.ctx.Done(): - return p.ctx.Err() - case p.idleConnections <- conn: - } - } + p.discard(conn) + + return p.stopped() + case p.idleConnections <- conn: } } } @@ -114,3 +137,55 @@ func (p *OneTimePool) Close() error { return p.InFlightConns.Close() } + +// stopped is the error Run returns once the pool context is done. +func (p *OneTimePool) stopped() error { + return fmt.Errorf("connpool: run stopped: %w", p.ctx.Err()) +} + +// dial opens one connection and returns nil if the dial failed; Run retries +// and does not report the cause, so the error is not built into a value. The +// dialer is read under the lock but invoked outside it, so Close never waits +// for a dial in progress. +func (p *OneTimePool) dial() net.Conn { + p.mx.RLock() + dialer := p.Dialer + p.mx.RUnlock() + + var ( + conn net.Conn + err error + ) + + if dialer != nil { + conn, err = dialer.Dial(p.ctx, "tcp", p.address) + } else { + var netDialer net.Dialer + + conn, err = netDialer.DialContext(p.ctx, "tcp", p.address) + } + + if err != nil { + return nil + } + + return conn +} + +// sleep waits for d and reports whether the pool context is still live. +func (p *OneTimePool) sleep(d time.Duration) bool { + timer := time.NewTimer(d) + defer timer.Stop() + + select { + case <-p.ctx.Done(): + return false + case <-timer.C: + return true + } +} + +func (p *OneTimePool) discard(conn net.Conn) { + p.InFlightConns.Remove(conn) + _ = conn.Close() +} diff --git a/conn/connpool/onetimepool_test.go b/conn/connpool/onetimepool_test.go index 78eb0bb..3d8400a 100644 --- a/conn/connpool/onetimepool_test.go +++ b/conn/connpool/onetimepool_test.go @@ -2,8 +2,13 @@ package connpool import ( "context" + "errors" + "io" "net" + "sync" + "sync/atomic" "testing" + "time" "github.com/stretchr/testify/require" ) @@ -27,3 +32,383 @@ func TestOneTimePoolAcquirePrefersCanceledContext(t *testing.T) { _ = pool.Close() } } + +// waitRunTimeout bounds how long a test waits for Run to return. +const waitRunTimeout = 5 * time.Second + +var errRefused = errors.New("connection refused") + +// funcDialer adapts a function to the Dialer interface. +type funcDialer func(ctx context.Context, network, address string) (net.Conn, error) + +func (f funcDialer) Dial(ctx context.Context, network, address string) (net.Conn, error) { + return f(ctx, network, address) +} + +// connTracker closes every connection handed to it when the test ends, so +// fake dialers and listeners never leak the peer side of a connection. +type connTracker struct { + mu sync.Mutex + conns []net.Conn + closed bool +} + +func newConnTracker(t *testing.T) *connTracker { + t.Helper() + + tracker := new(connTracker) + t.Cleanup(tracker.closeAll) + + return tracker +} + +func (c *connTracker) add(conn net.Conn) { + c.mu.Lock() + defer c.mu.Unlock() + + if c.closed { + _ = conn.Close() + + return + } + + c.conns = append(c.conns, conn) +} + +func (c *connTracker) first() net.Conn { + c.mu.Lock() + defer c.mu.Unlock() + + return c.conns[0] +} + +func (c *connTracker) closeAll() { + c.mu.Lock() + defer c.mu.Unlock() + + c.closed = true + for _, conn := range c.conns { + _ = conn.Close() + } +} + +// pipe returns the client side of a new in-memory connection and tracks the +// server side. +func (c *connTracker) pipe() net.Conn { + client, server := net.Pipe() + c.add(server) + + return client +} + +func newTestPool(t *testing.T, dialer Dialer, size int) *OneTimePool { + t.Helper() + + pool, err := NewOneTimePool(context.Background(), "unused", size, WithDialer(dialer)) + require.NoError(t, err) + t.Cleanup(func() { _ = pool.Close() }) + + return pool +} + +// runPool starts pool.Run in a goroutine and returns a channel that receives +// its result. +func runPool(pool *OneTimePool) <-chan error { + done := make(chan error, 1) + go func() { done <- pool.Run() }() + + return done +} + +func waitRun(t *testing.T, done <-chan error) error { + t.Helper() + + select { + case err := <-done: + return err + case <-time.After(waitRunTimeout): + t.Fatalf("Run did not return within %s", waitRunTimeout) + + return nil + } +} + +// waitAfterFailure is the wait Run applies after the nth consecutive failed +// dial (1-based), derived from the backoff constants. +func waitAfterFailure(n int) time.Duration { + wait := minDialBackoff + for range n - 1 { + wait = min(wait*backoffFactor, maxDialBackoff) + } + + return wait +} + +// maxDials is the most dial attempts Run can make in the given time when +// every dial fails. Timers never fire early, so attempt n+1 cannot start +// before the first n waits have elapsed. +func maxDials(elapsed time.Duration) int64 { + dials := int64(1) + + var waited time.Duration + + for n := 1; ; n++ { + waited += waitAfterFailure(n) + if waited > elapsed { + return dials + } + + dials++ + } +} + +func TestOneTimePoolRunBacksOffOnDialErrors(t *testing.T) { + t.Parallel() + + var dials atomic.Int64 + + pool := newTestPool(t, funcDialer(func(context.Context, string, string) (net.Conn, error) { + dials.Add(1) + + return nil, errRefused + }), 1) + + start := time.Now() + done := runPool(pool) + + time.Sleep(400 * time.Millisecond) + require.NoError(t, pool.Close()) + require.ErrorIs(t, waitRun(t, done), context.Canceled) + + // A tight retry loop with no wait between dials would blow far past the + // schedule-derived bound. + require.GreaterOrEqual(t, dials.Load(), int64(1)) + require.LessOrEqual(t, dials.Load(), maxDials(time.Since(start))) +} + +func TestOneTimePoolCloseInterruptsBackoff(t *testing.T) { + t.Parallel() + + // The wait after this many failures is the longest one still below the cap. + const failures = 5 + + var dials atomic.Int64 + + pool := newTestPool(t, funcDialer(func(context.Context, string, string) (net.Conn, error) { + dials.Add(1) + + return nil, errRefused + }), 1) + done := runPool(pool) + + require.Eventually(t, func() bool { return dials.Load() >= failures }, 10*time.Second, time.Millisecond) + + // Run is now parked in a wait of about pending. An uninterruptible wait + // would delay its return by most of that. + pending := waitAfterFailure(failures) + closeStart := time.Now() + + require.NoError(t, pool.Close()) + require.ErrorIs(t, waitRun(t, done), context.Canceled) + require.Less(t, time.Since(closeStart), pending/2) +} + +func TestOneTimePoolBackoffResetsAfterSuccess(t *testing.T) { + t.Parallel() + + // Failures before the one success. Without a reset, the wait after the + // next failure would be waitAfterFailure(failures+1). + const failures = 4 + + var ( + stampsMu sync.Mutex + stamps []time.Time + ) + + tracker := newConnTracker(t) + pool := newTestPool(t, funcDialer(func(context.Context, string, string) (net.Conn, error) { + stampsMu.Lock() + defer stampsMu.Unlock() + + stamps = append(stamps, time.Now()) + if len(stamps) == failures+1 { + return tracker.pipe(), nil + } + + return nil, errRefused + }), 1) + done := runPool(pool) + + // The pool holds one connection, so after the success Run dials again, + // fails, waits, and fails once more. + require.Eventually(t, func() bool { + stampsMu.Lock() + defer stampsMu.Unlock() + + return len(stamps) >= failures+3 + }, 10*time.Second, time.Millisecond) + require.NoError(t, pool.Close()) + require.ErrorIs(t, waitRun(t, done), context.Canceled) + + stampsMu.Lock() + defer stampsMu.Unlock() + + gap := stamps[failures+2].Sub(stamps[failures+1]) + require.GreaterOrEqual(t, gap, minDialBackoff) + require.Less(t, gap, waitAfterFailure(failures+1)/2) +} + +func TestOneTimePoolRecoversAfterDialErrors(t *testing.T) { + t.Parallel() + + var dials atomic.Int64 + + tracker := newConnTracker(t) + pool := newTestPool(t, funcDialer(func(context.Context, string, string) (net.Conn, error) { + if dials.Add(1) <= 2 { + return nil, errRefused + } + + return tracker.pipe(), nil + }), 1) + done := runPool(pool) + + ctx, cancel := context.WithTimeout(context.Background(), waitRunTimeout) + defer cancel() + + conn, err := pool.Acquire(ctx) + require.NoError(t, err) + require.NotNil(t, conn) + require.GreaterOrEqual(t, dials.Load(), int64(3)) + + go func() { _, _ = conn.Write([]byte("ping")) }() + + buf := make([]byte, len("ping")) + _, err = io.ReadFull(tracker.first(), buf) + require.NoError(t, err) + require.Equal(t, "ping", string(buf)) + + _ = conn.Close() + + require.NoError(t, pool.Close()) + require.ErrorIs(t, waitRun(t, done), context.Canceled) +} + +func TestOneTimePoolRunFillsPool(t *testing.T) { + t.Parallel() + + tracker := newConnTracker(t) + pool := newTestPool(t, funcDialer(func(context.Context, string, string) (net.Conn, error) { + return tracker.pipe(), nil + }), 2) + done := runPool(pool) + + ctx, cancel := context.WithTimeout(context.Background(), waitRunTimeout) + defer cancel() + + // More acquisitions than the pool size proves Run keeps refilling it. + for range 5 { + conn, err := pool.Acquire(ctx) + require.NoError(t, err) + require.NotNil(t, conn) + + _ = conn.Close() + } + + require.NoError(t, pool.Close()) + require.ErrorIs(t, waitRun(t, done), context.Canceled) +} + +func TestOneTimePoolRunWithoutDialer(t *testing.T) { + t.Parallel() + + var lc net.ListenConfig + + listener, err := lc.Listen(t.Context(), "tcp", "127.0.0.1:0") + require.NoError(t, err) + + tracker := newConnTracker(t) + t.Cleanup(func() { _ = listener.Close() }) + + go func() { + for { + accepted, err := listener.Accept() + if err != nil { + return + } + + tracker.add(accepted) + } + }() + + pool, err := NewOneTimePool(context.Background(), listener.Addr().String(), 1) + require.NoError(t, err) + t.Cleanup(func() { _ = pool.Close() }) + + done := runPool(pool) + + ctx, cancel := context.WithTimeout(context.Background(), waitRunTimeout) + defer cancel() + + conn, err := pool.Acquire(ctx) + require.NoError(t, err) + require.NotNil(t, conn) + + _ = conn.Close() + + require.NoError(t, pool.Close()) + require.ErrorIs(t, waitRun(t, done), context.Canceled) +} + +// Close must not wait for a dial in progress, and a connection that finishes +// dialing after Close must not be left open. +func TestOneTimePoolCloseDuringDial(t *testing.T) { + t.Parallel() + + var ( + entered = make(chan struct{}) + release = make(chan struct{}) + once sync.Once + ) + + tracker := newConnTracker(t) + dialed := tracker.pipe() + server := tracker.first() + pool := newTestPool(t, funcDialer(func(context.Context, string, string) (net.Conn, error) { + once.Do(func() { close(entered) }) + <-release + + return dialed, nil + }), 1) + done := runPool(pool) + + <-entered + + closed := make(chan error, 1) + go func() { closed <- pool.Close() }() + + select { + case err := <-closed: + require.NoError(t, err) + case <-time.After(time.Second): + close(release) + t.Fatal("Close blocked while a dial was in progress") + } + + close(release) + require.ErrorIs(t, waitRun(t, done), context.Canceled) + + // The dialed connection was discarded, so the peer observes EOF. + readErr := make(chan error, 1) + go func() { + _, err := server.Read(make([]byte, 1)) + readErr <- err + }() + + select { + case err := <-readErr: + require.ErrorIs(t, err, io.EOF) + case <-time.After(waitRunTimeout): + t.Fatal("dialed connection was left open after Close") + } +}