poll: interrupt blocked map writes

Updates #3410

hscontrol/poll.go

@@ -95,6 +95,15 @@ func (m *mapSession) resetKeepAlive() {
 func (m *mapSession) stopFromBatcher() {
 	if m.cancelChClosed.CompareAndSwap(false, true) {
 		close(m.cancelCh)
+
+		// A channel signal cannot interrupt a response write that is blocked on
+		// HTTP/2 flow control. Expire the stream's write deadline as well so a
+		// client that stopped reading cannot keep the map session, and therefore
+		// server shutdown, alive indefinitely.
+		err := http.NewResponseController(m.w).SetWriteDeadline(time.Now())
+		if err != nil && !errors.Is(err, http.ErrNotSupported) {
+			m.log.Debug().Caller().Err(err).Msg("failed to interrupt map response write")
+		}
 	}
 }
 

hscontrol/poll_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"net/http"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -33,6 +34,69 @@ type delayedSuccessResponseWriter struct {
 	writeCount int
 }
 
+// deadlineResponseWriter simulates an HTTP/2 response blocked on flow control.
+// The write only returns when the handler expires its per-stream write deadline.
+type deadlineResponseWriter struct {
+	header http.Header
+
+	blockWrites atomic.Bool
+	writeCount  atomic.Int64
+
+	writeStarted     chan struct{}
+	writeStartedOnce sync.Once
+	deadlineSet      chan struct{}
+	deadlineSetOnce  sync.Once
+}
+
+func newDeadlineResponseWriter() *deadlineResponseWriter {
+	return &deadlineResponseWriter{
+		header:       make(http.Header),
+		writeStarted: make(chan struct{}),
+		deadlineSet:  make(chan struct{}),
+	}
+}
+
+func (w *deadlineResponseWriter) Header() http.Header {
+	return w.header
+}
+
+func (w *deadlineResponseWriter) WriteHeader(int) {}
+
+func (w *deadlineResponseWriter) Write(data []byte) (int, error) {
+	w.writeCount.Add(1)
+
+	if !w.blockWrites.Load() {
+		return len(data), nil
+	}
+
+	w.writeStartedOnce.Do(func() { close(w.writeStarted) })
+	<-w.deadlineSet
+
+	return 0, context.DeadlineExceeded
+}
+
+func (w *deadlineResponseWriter) Flush() {}
+
+func (w *deadlineResponseWriter) SetWriteDeadline(deadline time.Time) error {
+	if !deadline.IsZero() {
+		w.deadlineSetOnce.Do(func() { close(w.deadlineSet) })
+	}
+
+	return nil
+}
+
+func (w *deadlineResponseWriter) WriteStarted() <-chan struct{} {
+	return w.writeStarted
+}
+
+func (w *deadlineResponseWriter) BlockWrites() {
+	w.blockWrites.Store(true)
+}
+
+func (w *deadlineResponseWriter) WriteCount() int64 {
+	return w.writeCount.Load()
+}
+
 func newDelayedSuccessResponseWriter(firstWriteDelay time.Duration) *delayedSuccessResponseWriter {
 	return &delayedSuccessResponseWriter{
 		header:             make(http.Header),
@@ -475,3 +539,69 @@ func TestDeletedNodeEndsLongPoll(t *testing.T) {
 		t.Fatal("an orphaned session would block server shutdown on clientStreamsOpen")
 	}
 }
+
+func TestDeletedNodeInterruptsBlockedWrite(t *testing.T) {
+	t.Parallel()
+
+	app := createTestApp(t)
+	user := app.state.CreateUserForTest("poll-delete-blocked-user")
+	createdNode := app.state.CreateRegisteredNodeForTest(user, "poll-delete-blocked-node")
+	require.NoError(t, app.state.UpdatePolicyManagerUsersForTest())
+	app.cfg.Tuning.BatchChangeDelay = 20 * time.Millisecond
+	app.cfg.Tuning.NodeMapSessionBufferedChanSize = 1
+
+	app.mapBatcher.Close()
+	require.NoError(t, app.state.Close())
+
+	reloadedState, err := state.NewState(app.cfg)
+	require.NoError(t, err)
+
+	app.state = reloadedState
+	app.mapBatcher = mapper.NewBatcherAndMapper(app.cfg, app.state)
+	app.mapBatcher.Start()
+
+	t.Cleanup(func() {
+		app.mapBatcher.Close()
+		require.NoError(t, app.state.Close())
+	})
+
+	nodeView, ok := app.state.GetNodeByID(createdNode.ID)
+	require.True(t, ok)
+
+	writer := newDeadlineResponseWriter()
+	session := app.newMapSession(t.Context(), tailcfg.MapRequest{
+		Stream:  true,
+		Version: tailcfg.CapabilityVersion(100),
+	}, writer, nodeView.AsStruct())
+
+	serveDone := make(chan struct{})
+
+	go func() {
+		session.serveLongPoll()
+		close(serveDone)
+	}()
+
+	require.Eventually(t, func() bool {
+		return writer.WriteCount() >= 2 && app.mapBatcher.IsConnected(nodeView.ID())
+	}, 2*time.Second, 10*time.Millisecond,
+		"expected the initial map and connect update to be delivered")
+
+	writer.BlockWrites()
+	app.Change(change.SelfUpdate(nodeView.ID()))
+
+	select {
+	case <-writer.WriteStarted():
+	case <-time.After(2 * time.Second):
+		t.Fatal("expected the initial map write to block")
+	}
+
+	c, err := app.state.DeleteNode(nodeView)
+	require.NoError(t, err)
+	app.Change(c)
+
+	select {
+	case <-serveDone:
+	case <-time.After(2 * time.Second):
+		t.Fatal("deleting a node must interrupt its blocked map response write")
+	}
+}