drpcmanager: add soft cancel grace period

drpcmanager/manager.go

@@ -26,6 +26,31 @@ import (
 
 var managerClosed = errs.Class("manager closed")
 
+// Outcomes reported to the cancel callback when a soft cancel is attempted. They
+// distinguish the case that keeps the transport from the two that destroy it,
+// which is the difference between reusing a connection and dialing a new one.
+const (
+	// CancelClean means the cancel was delivered and the transport was kept.
+	CancelClean = "clean"
+
+	// CancelBusy means the stream was still sending when the cancel arrived, so
+	// the transport had to be destroyed. If SoftCancelGrace is set, the grace
+	// period elapsed without the send finishing.
+	CancelBusy = "busy"
+
+	// CancelError means sending the cancel failed, so the transport had to be
+	// destroyed.
+	CancelError = "error"
+)
+
+// Bounds on how often a busy soft cancel retries within its grace period. The
+// first retry is cheap so that the common case of a send that is nearly done is
+// picked up quickly, and the backoff keeps a long grace from spinning.
+const (
+	softCancelMinBackoff = 100 * time.Microsecond
+	softCancelMaxBackoff = 5 * time.Millisecond
+)
+
 // Options controls configuration settings for a manager.
 type Options struct {
 	// WriterBufferSize controls the size of the buffer that we will fill before
@@ -45,6 +70,22 @@ type Options struct {
 	// being flushed when the cancel happens.
 	SoftCancel bool
 
+	// SoftCancelGrace bounds how long a soft cancel will wait for a concurrent
+	// send on the stream to finish before giving up and hard canceling. It only
+	// applies when SoftCancel is set.
+	//
+	// A stream that is sending when the cancel arrives cannot be softly canceled
+	// immediately, and without a grace period the transport is destroyed. For
+	// clients whose streams send and receive concurrently, that race is both
+	// benign and common, so paying a dial and a handshake for it is a poor
+	// trade. A grace period on the order of milliseconds lets the send finish so
+	// the cancel can be delivered and the transport reused.
+	//
+	// A wedged transport still gets torn down: the wait is bounded, and it never
+	// blocks on the write itself. If zero or negative, no grace is given and a
+	// concurrent send hard cancels immediately.
+	SoftCancelGrace time.Duration
+
 	// InactivityTimeout is the amount of time the manager will wait when
 	// creating a NewServerStream. It only includes the time it is reading
 	// packets from the remote client. In other words, it only includes the time
@@ -328,6 +369,73 @@ func (m *Manager) manageStreams() {
 	}
 }
 
+// cancelOutcome reports how a soft cancel attempt ended to the callback set in
+// the internal options, if any. It exists so that the ratio of kept to destroyed
+// transports is observable in production, where it is the difference between a
+// connection pool that works and one that never gets a hit.
+func (m *Manager) cancelOutcome(busy bool, err error) {
+	cb := drpcopts.GetManagerCancelCB(&m.opts.Internal)
+	if cb == nil {
+		return
+	}
+
+	switch {
+	case err != nil:
+		cb(CancelError)
+	case busy:
+		cb(CancelBusy)
+	default:
+		cb(CancelClean)
+	}
+}
+
+// sendCancel attempts to send a soft cancel for the stream, retrying within the
+// SoftCancelGrace window if the stream is busy sending something else.
+//
+// SendCancel never blocks: it reports busy rather than waiting on a mutex that a
+// wedged transport write may hold forever. That property is what lets a dead
+// transport still be torn down, so it is preserved here by polling with backoff
+// instead of waiting. All this adds is patience, bounded by the grace period,
+// for the far more common case of a send that is simply still in flight.
+func (m *Manager) sendCancel(ctx context.Context, stream *drpcstream.Stream) (busy bool, err error) {
+	defer func() { m.cancelOutcome(busy, err) }()
+
+	busy, err = stream.SendCancel(ctx.Err())
+	if err != nil || !busy || m.opts.SoftCancelGrace <= 0 {
+		return busy, err
+	}
+
+	grace := time.NewTimer(m.opts.SoftCancelGrace)
+	defer grace.Stop()
+
+	delay := softCancelMinBackoff
+	retry := time.NewTimer(delay)
+	defer retry.Stop()
+
+	for {
+		select {
+		case <-grace.C:
+			return true, nil
+
+		case <-m.sigs.term.Signal():
+			// the transport is already going away, so there is nothing left to
+			// preserve by waiting.
+			return true, nil
+
+		case <-retry.C:
+		}
+
+		if busy, err = stream.SendCancel(ctx.Err()); err != nil || !busy {
+			return busy, err
+		}
+
+		if delay *= 2; delay > softCancelMaxBackoff {
+			delay = softCancelMaxBackoff
+		}
+		retry.Reset(delay)
+	}
+}
+
 // manageStream watches the context and the stream and returns when the stream
 // is finished, canceling the stream if the context is canceled.
 func (m *Manager) manageStream(ctx context.Context, stream *drpcstream.Stream) {
@@ -353,7 +461,7 @@ func (m *Manager) manageStream(ctx context.Context, stream *drpcstream.Stream) {
 
 			// attempt to send the soft cancel. if it fails or if the stream is
 			// busy sending something else, then we have to hard cancel.
-			if busy, err := stream.SendCancel(ctx.Err()); err != nil {
+			if busy, err := m.sendCancel(ctx, stream); err != nil {
 				m.terminate(err)
 			} else if busy {
 				m.log("BUSY", stream.String)

drpcmanager/manager_test.go

@@ -16,6 +16,7 @@ import (
 
 	"storj.io/drpc/drpctest"
 	"storj.io/drpc/drpcwire"
+	"storj.io/drpc/internal/drpcopts"
 )
 
 func closed(ch <-chan struct{}) bool {
@@ -127,11 +128,81 @@ func TestUnblocked_SoftCancel(t *testing.T) {
 	t.Run("Disabled", func(t *testing.T) { run(t, false) })
 }
 
+// TestSoftCancel_Grace covers the case that a stream is sending when its context
+// is canceled. Without a grace period the transport is destroyed; with one, the
+// send is given a bounded chance to finish so the cancel can be delivered and the
+// transport kept.
+func TestSoftCancel_Grace(t *testing.T) {
+	run := func(t *testing.T, grace time.Duration) string {
+		ctx := drpctest.NewTracker(t)
+		defer ctx.Close()
+
+		outcomes := make(chan string, 1)
+
+		tr := newBlockedTransport()
+		opts := Options{SoftCancel: true, SoftCancelGrace: grace}
+		drpcopts.SetManagerCancelCB(&opts.Internal, func(o string) {
+			select {
+			case outcomes <- o:
+			default:
+			}
+		})
+
+		man := NewWithOptions(tr, opts)
+		defer func() { _ = man.Close() }()
+		defer tr.setReadOpen(true)
+		defer tr.setWriteOpen(true)
+
+		subctx, cancel := context.WithCancel(ctx)
+		defer cancel()
+
+		stream, err := man.NewClientStream(subctx, "rpc")
+		assert.NoError(t, err)
+		defer func() { _ = stream.Close() }()
+
+		// hold the stream's write mutex the way a concurrent send would, by
+		// parking an in-flight flush inside the transport.
+		ctx.Run(func(context.Context) {
+			_ = stream.RawWrite(drpcwire.KindMessage, []byte("message"))
+			_ = stream.RawFlush()
+		})
+
+		// wait until the flush is actually parked in the transport, so the cancel
+		// below reliably observes the stream as busy.
+		tr.waitWriting()
+
+		// release the send only after the manager has had a chance to see it as
+		// busy. Without a grace period the manager gives up before this fires and
+		// reports busy; with one, the retry picks the stream up once it lands.
+		time.AfterFunc(50*time.Millisecond, func() { tr.setWriteOpen(true) })
+
+		cancel()
+
+		select {
+		case outcome := <-outcomes:
+			return outcome
+		case <-time.After(10 * time.Second):
+			t.Fatal("timed out waiting for a cancel outcome")
+			return ""
+		}
+	}
+
+	t.Run("NoGrace", func(t *testing.T) {
+		assert.Equal(t, run(t, 0), CancelBusy)
+	})
+
+	t.Run("Grace", func(t *testing.T) {
+		assert.Equal(t, run(t, time.Second), CancelClean)
+	})
+}
+
 type blockedTransport struct {
 	mu *sync.Mutex
 	co *sync.Cond
 	ro bool
 	wo bool
+	rn int // number of reads currently blocked
+	wn int // number of writes currently blocked
 }
 
 func newBlockedTransport() *blockedTransport {
@@ -159,16 +230,30 @@ func (b *blockedTransport) setReadOpen(open bool) {
 	b.co.Broadcast()
 }
 
-func (b *blockedTransport) wait(p int, rw *bool) (int, error) {
+func (b *blockedTransport) wait(p int, rw *bool, n *int) (int, error) {
 	b.mu.Lock()
 	defer b.mu.Unlock()
 
+	*n++
+	b.co.Broadcast()
+	defer func() { *n-- }()
+
 	for !*rw {
 		b.co.Wait()
 	}
 	return p, nil
 }
 
-func (b *blockedTransport) Read(p []byte) (n int, err error)  { return b.wait(len(p), &b.ro) }
-func (b *blockedTransport) Write(p []byte) (n int, err error) { return b.wait(len(p), &b.wo) }
+// waitWriting blocks until at least one Write is parked in the transport.
+func (b *blockedTransport) waitWriting() {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	for b.wn == 0 {
+		b.co.Wait()
+	}
+}
+
+func (b *blockedTransport) Read(p []byte) (n int, err error)  { return b.wait(len(p), &b.ro, &b.rn) }
+func (b *blockedTransport) Write(p []byte) (n int, err error) { return b.wait(len(p), &b.wo, &b.wn) }
 func (b *blockedTransport) Close() error                      { return nil }

internal/drpcopts/manager.go

@@ -7,11 +7,18 @@ import "storj.io/drpc/drpcstats"
 
 // Manager contains internal options for the drpcmanager package.
 type Manager struct {
-	statsCB func(string) *drpcstats.Stats
+	statsCB  func(string) *drpcstats.Stats
+	cancelCB func(string)
 }
 
 // GetManagerStatsCB returns the stats callback stored in the options.
 func GetManagerStatsCB(opts *Manager) func(string) *drpcstats.Stats { return opts.statsCB }
 
 // SetManagerStatsCB sets the stats callback stored in the options.
 func SetManagerStatsCB(opts *Manager, statsCB func(string) *drpcstats.Stats) { opts.statsCB = statsCB }
+
+// GetManagerCancelCB returns the cancel outcome callback stored in the options.
+func GetManagerCancelCB(opts *Manager) func(string) { return opts.cancelCB }
+
+// SetManagerCancelCB sets the cancel outcome callback stored in the options.
+func SetManagerCancelCB(opts *Manager, cancelCB func(string)) { opts.cancelCB = cancelCB }