fix: bound native shutdown and close connections that become idle during drain

transport/native/server.go

@@ -41,6 +41,7 @@ const (
 // It supports keep-alive, graceful shutdown, bounded concurrency and standard
 // http.Handler values. Zero MaxConnections and MaxConcurrentRequests mean
 // unlimited; pooled protocol objects are reused automatically.
+// A Server owns one serving lifecycle; create a new Server to restart.
 type Server struct {
 	ReadHeaderTimeout     time.Duration
 	ReadTimeout           time.Duration
@@ -232,9 +233,9 @@ func (s *Server) Serve(ctx context.Context, listener net.Listener) error {
 	}
 
 	s.mu.Lock()
-	if s.listener != nil {
+	if s.cancel != nil {
 		s.mu.Unlock()
-		return errors.New("go-omni/native: server is already running")
+		return errors.New("go-omni/native: server has already been started")
 	}
 	s.listener = listener
 	s.closing = false
@@ -357,7 +358,8 @@ func (s *Server) Shutdown(ctx context.Context) error {
 		for _, state := range states {
 			state.close()
 		}
-		<-done
+		// Handlers must cooperate with cancellation. Waiting here would make
+		// an uncooperative handler defeat the caller's shutdown deadline.
 		return ctx.Err()
 	}
 }
@@ -446,8 +448,7 @@ func (s *Server) serveConnection(baseCtx context.Context, state *connection) {
 				cancel()
 				writer.reset("")
 				s.writerPool.Put(writer)
-				state.active.Store(false)
-				s.activeRequests.Add(-1)
+				s.finishRequest(state)
 				return
 			default:
 				s.rejectedRequests.Add(1)
@@ -459,8 +460,7 @@ func (s *Server) serveConnection(baseCtx context.Context, state *connection) {
 				cancel()
 				writer.reset("")
 				s.writerPool.Put(writer)
-				state.active.Store(false)
-				s.activeRequests.Add(-1)
+				closing = s.finishRequest(state) || closing
 				if closeAfter || closing {
 					return
 				}
@@ -496,8 +496,7 @@ func (s *Server) serveConnection(baseCtx context.Context, state *connection) {
 				}
 				writer.reset("")
 				s.writerPool.Put(writer)
-				state.active.Store(false)
-				s.activeRequests.Add(-1)
+				s.finishRequest(state)
 				return
 			}
 			// Request bodies are streamed into the handler, so a malformed
@@ -518,8 +517,7 @@ func (s *Server) serveConnection(baseCtx context.Context, state *connection) {
 			}
 			writer.reset("")
 			s.writerPool.Put(writer)
-			state.active.Store(false)
-			s.activeRequests.Add(-1)
+			s.finishRequest(state)
 			return
 		}
 		if requestSlot {
@@ -528,14 +526,25 @@ func (s *Server) serveConnection(baseCtx context.Context, state *connection) {
 		shouldClose := closeAfter || writer.closeAfter || closing
 		writer.reset("")
 		s.writerPool.Put(writer)
-		state.active.Store(false)
-		s.activeRequests.Add(-1)
+		shouldClose = s.finishRequest(state) || shouldClose
 		if shouldClose {
 			return
 		}
 	}
 }
 
+// Publish the idle transition under the same lock that starts shutdown. Either
+// Shutdown observes an idle connection and closes it, or this request observes
+// shutdown and closes it instead of starting another Keep-Alive read.
+func (s *Server) finishRequest(state *connection) bool {
+	s.mu.Lock()
+	state.active.Store(false)
+	closing := s.closing || state.closeWhenIdle.Load()
+	s.mu.Unlock()
+	s.activeRequests.Add(-1)
+	return closing
+}
+
 func (c *connection) close() {
 	if c.closed.CompareAndSwap(false, true) {
 		defer close(c.closeDone)

transport/native/shutdown_test.go

@@ -0,0 +1,101 @@
+package native
+
+import (
+	"bufio"
+	"context"
+	"io"
+	"net"
+	"net/http"
+	"sync"
+	"testing"
+	"time"
+)
+
+type pausedWriteListener struct {
+	net.Listener
+	writing chan struct{}
+	release chan struct{}
+}
+
+func (l *pausedWriteListener) Accept() (net.Conn, error) {
+	conn, err := l.Listener.Accept()
+	if err != nil {
+		return nil, err
+	}
+	return &pausedWriteConn{Conn: conn, writing: l.writing, release: l.release}, nil
+}
+
+type pausedWriteConn struct {
+	net.Conn
+	writing chan struct{}
+	release chan struct{}
+	once    sync.Once
+}
+
+func (c *pausedWriteConn) Write(data []byte) (int, error) {
+	c.once.Do(func() { close(c.writing); <-c.release })
+	return c.Conn.Write(data)
+}
+
+func TestShutdownDuringResponseWriteClosesNewlyIdleConnection(t *testing.T) {
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	wrapped := &pausedWriteListener{Listener: listener, writing: make(chan struct{}), release: make(chan struct{})}
+	var once sync.Once
+	unblock := func() { once.Do(func() { close(wrapped.release) }) }
+	defer unblock()
+	defer listener.Close()
+	server := New(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { _, _ = io.WriteString(w, "finished") }))
+	server.ShutdownTimeout = 200 * time.Millisecond
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	done := make(chan error, 1)
+	go func() { done <- server.Serve(ctx, wrapped) }()
+	client, err := net.DialTimeout("tcp", listener.Addr().String(), time.Second)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer client.Close()
+	client.SetDeadline(time.Now().Add(2 * time.Second))
+	if _, err := io.WriteString(client, "GET / HTTP/1.1\r\nHost: localhost\r\n\r\n"); err != nil {
+		t.Fatal(err)
+	}
+	select {
+	case <-wrapped.writing:
+	case <-time.After(time.Second):
+		t.Fatal("response did not reach write")
+	}
+	// The response has already chosen its Keep-Alive header. Shutdown must
+	// still close this connection when its write completes and it becomes idle.
+	cancel()
+	deadline := time.Now().Add(time.Second)
+	for !server.isClosing() {
+		if time.Now().After(deadline) {
+			t.Fatal("shutdown did not start")
+		}
+		time.Sleep(time.Millisecond)
+	}
+	unblock()
+	response, err := http.ReadResponse(bufio.NewReader(client), &http.Request{Method: "GET"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	body, err := io.ReadAll(response.Body)
+	response.Body.Close()
+	if err != nil || string(body) != "finished" {
+		t.Fatalf("response %q: %v", body, err)
+	}
+	select {
+	case err := <-done:
+		if err != nil {
+			t.Fatalf("completed request left a connection stuck in drain: %v", err)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("shutdown did not complete")
+	}
+	if stats := server.Stats(); stats.ActiveConnections != 0 || stats.ActiveRequests != 0 {
+		t.Fatalf("resources retained: %+v", stats)
+	}
+}