feat: cancel native requests on disconnect

transport/native/server_test.go

@@ -496,6 +496,121 @@ func TestServerShutdownCancelsActiveHandler(t *testing.T) {
 	}
 }
 
+func TestServerCancelsHandlerWhenClientDisconnects(t *testing.T) {
+	for _, stream := range []bool{false, true} {
+		t.Run(fmt.Sprintf("stream=%t", stream), func(t *testing.T) {
+			started := make(chan struct{})
+			canceled := make(chan struct{})
+			application := omni.NewWithRouter(netrouter.New())
+			application.GET("/wait", func(ctx *omni.Context) error {
+				if stream {
+					_, _ = io.WriteString(ctx.Raw(), "ready\n")
+					ctx.Raw().(http.Flusher).Flush()
+				}
+				close(started)
+				<-ctx.Context().Done()
+				close(canceled)
+				return nil
+			})
+			server, listener, done := startServer(t, application)
+			defer stopServer(t, server, listener, done)
+
+			conn, err := net.Dial("tcp", listener.Addr().String())
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer conn.Close()
+			conn.SetDeadline(time.Now().Add(2 * time.Second))
+			_, _ = io.WriteString(conn, "GET /wait HTTP/1.1\r\nHost: test\r\nConnection: close\r\n\r\n")
+			select {
+			case <-started:
+			case <-time.After(time.Second):
+				t.Fatal("handler did not start")
+			}
+			if stream {
+				response, err := http.ReadResponse(bufio.NewReader(conn), &http.Request{Method: http.MethodGet})
+				if err != nil {
+					t.Fatal(err)
+				}
+				first := make([]byte, len("ready\n"))
+				if _, err := io.ReadFull(response.Body, first); err != nil || string(first) != "ready\n" {
+					t.Fatalf("failed reading stream: %q err=%v", first, err)
+				}
+				// Close the socket before closing the unfinished response body.
+				conn.Close()
+				response.Body.Close()
+			} else {
+				conn.Close()
+			}
+			select {
+			case <-canceled:
+			case <-time.After(time.Second):
+				t.Fatal("client disconnect did not cancel request context")
+			}
+		})
+	}
+}
+
+func TestServerLongHandlerPreservesKeepAliveAndPipeline(t *testing.T) {
+	for _, pipeline := range []bool{false, true} {
+		t.Run(fmt.Sprintf("pipeline=%t", pipeline), func(t *testing.T) {
+			application := omni.New()
+			application.GET("/slow", func(ctx *omni.Context) error {
+				timer := time.NewTimer(25 * time.Millisecond)
+				defer timer.Stop()
+				select {
+				case <-timer.C:
+					return ctx.OK("slow done")
+				case <-ctx.Context().Done():
+					return ctx.Context().Err()
+				}
+			})
+			application.POST("/next", func(ctx *omni.Context) error {
+				body, err := io.ReadAll(ctx.Request().Body)
+				if err != nil {
+					return err
+				}
+				return ctx.OK(string(body))
+			})
+			server, listener, done := startServer(t, application)
+			defer stopServer(t, server, listener, done)
+			conn, err := net.Dial("tcp", listener.Addr().String())
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer conn.Close()
+			conn.SetDeadline(time.Now().Add(3 * time.Second))
+			reader := bufio.NewReader(conn)
+			for attempt := 0; attempt < 3; attempt++ {
+				slow := "GET /slow HTTP/1.1\r\nHost: test\r\n\r\n"
+				next := "POST /next HTTP/1.1\r\nHost: test\r\nContent-Length: 7\r\n\r\npayload"
+				if pipeline {
+					slow += next
+				}
+				if _, err := io.WriteString(conn, slow); err != nil {
+					t.Fatal(err)
+				}
+				for index, want := range []string{"slow done", "payload"} {
+					if index == 1 && !pipeline {
+						if _, err := io.WriteString(conn, next); err != nil {
+							t.Fatal(err)
+						}
+					}
+					response, err := http.ReadResponse(reader, &http.Request{Method: http.MethodGet})
+					if err != nil {
+						t.Fatal(err)
+					}
+					body, err := io.ReadAll(response.Body)
+					response.Body.Close()
+					if err != nil || response.StatusCode != http.StatusOK || !strings.Contains(string(body), want) {
+						t.Fatalf("attempt=%d response=%d status=%d body=%s err=%v", attempt, index, response.StatusCode, body, err)
+					}
+				}
+			}
+		})
+	}
+}
+
 func TestServerConcurrentRequestLimit(t *testing.T) {
 	started := make(chan struct{})
 	release := make(chan struct{})