poll: let a stopped session's write in flight finish

Expiring the deadline at once still reset streams whose last write
was in flight; a grace keeps the bound on clients that stopped reading.

hscontrol/poll.go

@@ -24,6 +24,10 @@ 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
@@ -107,15 +111,17 @@ func (m *mapSession) stopFromBatcher() {
 		// A channel signal cannot interrupt a response write that is blocked on
 		// HTTP/2 flow control, so a client that stopped reading would keep the
 		// map session, and therefore server shutdown, alive indefinitely.
-		// Expiring the write deadline resets the stream with INTERNAL_ERROR,
+		// An expired write deadline resets the stream with INTERNAL_ERROR,
-		// so only do it to a write in flight; an idle session sees cancelCh
+		// which the client reports as a server error, so it is only a bound:
-		// and ends the stream cleanly. writeMap checks cancelChClosed after
+		// an idle session sees cancelCh, and a write the client keeps reading
-		// raising writing, so one side always sees the other.
+		// 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())
+		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")
 		}

hscontrol/poll_test.go

@@ -5,6 +5,7 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"testing"
@@ -664,13 +665,27 @@ func TestDeletedNodeEndsLongPoll(t *testing.T) {
 }
 
 // TestDeletedNodeEndsStreamCleanly serves a long poll over HTTP/2, as the
-// Noise server does, and deletes the node while the session is idle. Nothing
+// Noise server does, and deletes the node, a clean stop, so the stream must
-// is left to interrupt, so the stream must end with END_STREAM: an expired
+// end with END_STREAM: an expired write deadline resets it with
-// write deadline resets it with INTERNAL_ERROR, which the client reports as a
+// INTERNAL_ERROR, which the client reports as a server error. That holds
-// server error.
+// 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
 
@@ -691,11 +706,17 @@ func TestDeletedNodeEndsStreamCleanly(t *testing.T) {
 			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) {
-		app.newMapSession(r.Context(), tailcfg.MapRequest{
+				session := app.newMapSession(r.Context(), tailcfg.MapRequest{
 					Stream:  true,
 					Version: tailcfg.CapabilityVersion(100),
-		}, w, nodeView.AsStruct()).serveLongPoll()
+				}, w, nodeView.AsStruct())
+				sessions <- session
+
+				session.serveLongPoll()
+				close(served)
 			}))
 			srv.Config.Protocols = new(http.Protocols)
 			srv.Config.Protocols.SetUnencryptedHTTP2(true)
@@ -704,7 +725,11 @@ func TestDeletedNodeEndsStreamCleanly(t *testing.T) {
 
 			h2c := new(http.Protocols)
 			h2c.SetUnencryptedHTTP2(true)
-	client := &http.Client{Transport: &http.Transport{Protocols: h2c}}
+			// 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)
@@ -719,12 +744,32 @@ func TestDeletedNodeEndsStreamCleanly(t *testing.T) {
 			_, 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, "stopping an idle session must end the stream, not reset it")
+			require.NoError(t, err, "a clean stop must end the stream, not reset it")
+		})
+	}
 }
 
 func TestDeletedNodeInterruptsBlockedWrite(t *testing.T) {