poll: end stopped map sessions without resetting active streams

Backport 371674de9 and 73dd39982 from #3575.

hscontrol/poll.go

@@ -24,8 +24,16 @@ import (
 
 const (
 	keepAliveInterval = 50 * time.Second
+
+	// stopWriteGrace is how long a stopped session's write in flight may
+	// take to finish before its stream is reset.
+	stopWriteGrace = 5 * time.Second
 )
 
+// errMapSessionStopped is a write refused because the batcher stopped the
+// session; the poll loop ends as if it had seen cancelCh.
+var errMapSessionStopped = errors.New("map session stopped")
+
 type contextKey string
 
 const nodeNameContextKey = contextKey("nodeName")
@@ -40,6 +48,10 @@ type mapSession struct {
 	cancelCh       chan struct{}
 	cancelChClosed atomic.Bool
 
+	// writing is set while [mapSession.writeMap] runs, so
+	// [mapSession.stopFromBatcher] interrupts only a write that may be stuck.
+	writing atomic.Bool
+
 	keepAlive       time.Duration
 	keepAliveTicker *time.Ticker
 
@@ -97,10 +109,19 @@ func (m *mapSession) stopFromBatcher() {
 		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())
+		// HTTP/2 flow control, so a client that stopped reading would keep the
+		// map session, and therefore server shutdown, alive indefinitely.
+		// An expired write deadline resets the stream with INTERNAL_ERROR,
+		// which the client reports as a server error, so it is only a bound:
+		// an idle session sees cancelCh, and a write the client keeps reading
+		// finishes first, both ending the stream cleanly. writeMap checks
+		// cancelChClosed after raising writing, so one side always sees the
+		// other.
+		if !m.writing.Load() {
+			return
+		}
+
+		err := http.NewResponseController(m.w).SetWriteDeadline(time.Now().Add(stopWriteGrace))
 		if err != nil && !errors.Is(err, http.ErrNotSupported) {
 			m.log.Debug().Caller().Err(err).Msg("failed to interrupt map response write")
 		}
@@ -318,6 +339,12 @@ func (m *mapSession) serveLongPoll() {
 			}
 
 			err := m.writeMap(update)
+			if errors.Is(err, errMapSessionStopped) {
+				mapResponseEnded.WithLabelValues("cancelled").Inc()
+
+				return
+			}
+
 			if err != nil {
 				m.log.Error().Caller().Err(err).Msg("cannot write update to client")
 				return
@@ -328,6 +355,12 @@ func (m *mapSession) serveLongPoll() {
 
 		case <-m.keepAliveTicker.C:
 			err := m.writeMap(&keepAlive)
+			if errors.Is(err, errMapSessionStopped) {
+				mapResponseEnded.WithLabelValues("cancelled").Inc()
+
+				return
+			}
+
 			if err != nil {
 				m.log.Error().Caller().Err(err).Msg("cannot write keep alive")
 				return
@@ -348,6 +381,13 @@ func (m *mapSession) serveLongPoll() {
 // It also handles flushing the response if the [http.ResponseWriter]
 // implements [http.Flusher].
 func (m *mapSession) writeMap(msg *tailcfg.MapResponse) error {
+	m.writing.Store(true)
+	defer m.writing.Store(false)
+
+	if m.cancelChClosed.Load() {
+		return errMapSessionStopped
+	}
+
 	startWrite := time.Now()
 
 	err := writeMapResponse(m.w, m.req.Compress, m.isStreaming(), msg)

hscontrol/poll_test.go

@@ -2,7 +2,10 @@ package hscontrol
 
 import (
 	"context"
+	"io"
 	"net/http"
+	"net/http/httptest"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"testing"
@@ -622,6 +625,114 @@ func TestDeletedNodeEndsLongPoll(t *testing.T) {
 	}
 }
 
+// TestDeletedNodeEndsStreamCleanly serves a long poll over HTTP/2, as the
+// Noise server does, and deletes the node, a clean stop, so the stream must
+// end with END_STREAM: an expired write deadline resets it with
+// INTERNAL_ERROR, which the client reports as a server error. That holds
+// with a write in flight too, as long as the client keeps reading; a client
+// that stopped reading must still not keep the session past the grace.
+func TestDeletedNodeEndsStreamCleanly(t *testing.T) {
+	t.Parallel()
+
+	for _, tt := range []struct {
+		name              string
+		inFlight, stalled bool
+	}{
+		{name: "idle"},
+		{name: "write_in_flight", inFlight: true},
+		{name: "client_stalled", inFlight: true, stalled: true},
+	} {
+		inFlight := tt.inFlight
+
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			app := createTestApp(t)
+			app.cfg.Tuning.NodeMapSessionBufferedChanSize = 30
+
+			user := app.state.CreateUserForTest("poll-delete-clean-user")
+			pak, err := app.state.CreatePreAuthKey(user.TypedID(), false, false, nil, nil)
+			require.NoError(t, err)
+
+			nodeKey := key.NewNode()
+
+			_, err = app.handleRegister(context.Background(), tailcfg.RegisterRequest{
+				Auth:     &tailcfg.RegisterResponseAuth{AuthKey: pak.Key},
+				NodeKey:  nodeKey.Public(),
+				Hostinfo: &tailcfg.Hostinfo{Hostname: "poll-delete-clean-node"},
+				Expiry:   time.Now().Add(24 * time.Hour),
+			}, key.NewMachine().Public())
+			require.NoError(t, err)
+
+			nodeView, ok := app.state.GetNodeByNodeKey(nodeKey.Public())
+			require.True(t, ok)
+
+			sessions := make(chan *mapSession, 1)
+			served := make(chan struct{})
+			srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				session := app.newMapSession(r.Context(), tailcfg.MapRequest{
+					Stream:  true,
+					Version: tailcfg.CapabilityVersion(100),
+				}, w, nodeView.AsStruct())
+				sessions <- session
+
+				session.serveLongPoll()
+				close(served)
+			}))
+			srv.Config.Protocols = new(http.Protocols)
+			srv.Config.Protocols.SetUnencryptedHTTP2(true)
+			srv.Start()
+			t.Cleanup(srv.Close)
+
+			h2c := new(http.Protocols)
+			h2c.SetUnencryptedHTTP2(true)
+			// A small window lets a client that pauses reading block a write.
+			client := &http.Client{Transport: &http.Transport{
+				Protocols: h2c,
+				HTTP2:     &http.HTTP2Config{MaxReceiveBufferPerStream: 64 << 10},
+			}}
+
+			req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, srv.URL, nil)
+			require.NoError(t, err)
+
+			resp, err := client.Do(req)
+			require.NoError(t, err)
+
+			defer resp.Body.Close()
+
+			var size [4]byte
+
+			_, err = io.ReadFull(resp.Body, size[:])
+			require.NoError(t, err, "expected the initial map")
+
+			if inFlight {
+				session := <-sessions
+				session.ch <- &tailcfg.MapResponse{Domain: strings.Repeat("x", 1<<20)}
+
+				require.Eventually(t, session.writing.Load, 2*time.Second, 10*time.Millisecond,
+					"test sanity: the frame must outgrow the window and block")
+			}
+
+			cs, err := app.state.DeleteNode(nodeView)
+			require.NoError(t, err)
+			app.Change(cs...)
+
+			if tt.stalled {
+				select {
+				case <-served:
+				case <-time.After(stopWriteGrace + 3*time.Second):
+					t.Fatal("a client that stopped reading must not keep a stopped session")
+				}
+
+				return
+			}
+
+			_, err = io.Copy(io.Discard, resp.Body)
+			require.NoError(t, err, "a clean stop must end the stream, not reset it")
+		})
+	}
+}
+
 func TestDeletedNodeInterruptsBlockedWrite(t *testing.T) {
 	t.Parallel()