From d36f0c603966a93f51cda13bc4cd5ee3473f5085 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?I=C3=B1aki=20Soria?=
 <77835286+inakisoriamrf@users.noreply.github.com>
Date: Thu, 1 Oct 2026 19:32:47 +0200
Subject: [PATCH] Do not fail awaiting cache transactions when the client
 closes

When a cacheable query is in flight, identical queries await its cache
transaction. If the client of the first query closed the connection,
the query was killed and the transaction was marked as failed. The
awaiting queries then got HTTP 500 "[concurrent query failed]" and did
not run, although nothing was wrong with the query.

Mark the transaction as completed when the query was cancelled because
its client closed the connection. The awaiting queries find no cached
result and run the query themselves. Timeouts and ClickHouse errors
still fail the transaction.
---
 proxy.go      |   7 +++-
 proxy_test.go | 112 ++++++++++++++++++++++++++++++++++++++++++++++++++
 scope.go      |   5 +++
 3 files changed, 122 insertions(+), 2 deletions(-)

diff --git a/proxy.go b/proxy.go
index 72ec90bb..f51aec4e 100644
--- a/proxy.go
+++ b/proxy.go
@@ -325,6 +325,7 @@ func (rp *reverseProxy) proxyRequest(s *scope, rw ResponseWriterWithCode, srw *s
 		return
 	case errors.Is(err, context.Canceled):
 		canceledRequest.With(s.labels).Inc()
+		s.clientClosed = true
 
 		q := getQuerySnippet(req)
 		log.Debugf("%s: remote client closed the connection in %s; query: %q", s, time.Since(startTime), q)
@@ -586,8 +587,10 @@ func (rp *reverseProxy) completeTransaction(s *scope, statusCode int, userCache
 	q []byte,
 	failReason string,
 ) {
-	// complete successful transactions or those with empty fail reason
-	if statusCode < 300 || failReason == "" {
+	// complete successful transactions or those with empty fail reason.
+	// A query cancelled because its client closed the connection did not fail:
+	// complete it too, so the concurrent queries that await it run the query themselves.
+	if statusCode < 300 || failReason == "" || s.clientClosed {
 		if err := userCache.Complete(key); err != nil {
 			log.Errorf("%s: %s; query: %q", s, err, q)
 		}
diff --git a/proxy_test.go b/proxy_test.go
index aa2f149a..0384d08d 100644
--- a/proxy_test.go
+++ b/proxy_test.go
@@ -871,6 +871,118 @@ func TestReverseProxy_ServeHTTP1(t *testing.T) {
 	})
 }
 
+// closeNotifierOnSignal reports that the client closed the connection when closed receives a value.
+type closeNotifierOnSignal struct {
+	http.ResponseWriter
+	closed chan bool
+}
+
+func (c *closeNotifierOnSignal) CloseNotify() <-chan bool { return c.closed }
+
+func waitFor(t *testing.T, what string, cond func() bool) {
+	t.Helper()
+	deadline := time.Now().Add(5 * time.Second)
+	for !cond() {
+		if time.Now().After(deadline) {
+			t.Fatalf("timed out waiting for %s", what)
+		}
+		time.Sleep(time.Millisecond)
+	}
+}
+
+func TestReverseProxy_ClientCloseCompletesCacheTransaction(t *testing.T) {
+	// With a grace time, identical queries await the cache transaction of the first one.
+	cfg := *goodCfgWithCache
+	cfg.Caches = append([]config.Cache(nil), goodCfgWithCache.Caches...)
+	cfg.Caches[0].GraceTime = config.Duration(5 * time.Second)
+
+	newRequest := func(query string) *http.Request {
+		return httptest.NewRequest("POST", fmt.Sprintf("%s?query=%s", fakeServer.URL, query), bytes.NewBufferString("1s"))
+	}
+	assertQueryRan := func(t *testing.T, resp *http.Response) {
+		t.Helper()
+		b := bbToString(t, resp.Body)
+		resp.Body.Close()
+		if resp.StatusCode != http.StatusOK {
+			t.Fatalf("unexpected status code: %d; expected: %d; response: %q", resp.StatusCode, http.StatusOK, b)
+		}
+		if strings.Contains(b, failedTransactionPrefix) {
+			t.Fatalf("the query got the error of a cancelled concurrent query: %q", b)
+		}
+	}
+
+	t.Run("query that awaits the cancelled one", func(t *testing.T) {
+		stopAllRequestsInFlight()
+		proxy, err := getProxy(&cfg)
+		if err != nil {
+			t.Fatalf("unexpected error: %s", err)
+		}
+		query := fmt.Sprintf("SELECT_client_close_await_%d", time.Now().UnixNano())
+
+		// A: a slow cacheable query whose client closes the connection on demand.
+		closeA := make(chan bool, 1)
+		doneA := make(chan struct{})
+		go func() {
+			defer close(doneA)
+			proxy.ServeHTTP(&closeNotifierOnSignal{ResponseWriter: httptest.NewRecorder(), closed: closeA}, newRequest(query))
+		}()
+		waitFor(t, "the first query to run on the server", func() bool {
+			return atomic.LoadInt64(&nbRequestsInflight) == 1
+		})
+
+		// B: the same query awaits the cache transaction of A.
+		respB := make(chan *http.Response, 1)
+		go func() { respB <- makeCustomRequest(proxy, newRequest(query)) }()
+		// Longer than one poll of AwaitForConcurrentTransaction (100 ms). B must still wait,
+		// and must not have reached the server: it awaits the transaction of A.
+		time.Sleep(150 * time.Millisecond)
+		select {
+		case resp := <-respB:
+			t.Fatalf("the second query did not await the first one; status code: %d", resp.StatusCode)
+		default:
+		}
+		if n := atomic.LoadInt64(&nbRequestsInflight); n != 1 {
+			t.Fatalf("expected only the first query on the server, got %d queries", n)
+		}
+
+		// The client of A goes away while B awaits.
+		closeA <- true
+		<-doneA
+
+		select {
+		case resp := <-respB:
+			assertQueryRan(t, resp)
+		case <-time.After(5 * time.Second):
+			t.Fatalf("the second query did not end")
+		}
+	})
+
+	t.Run("same query after the cancelled one", func(t *testing.T) {
+		stopAllRequestsInFlight()
+		proxy, err := getProxy(&cfg)
+		if err != nil {
+			t.Fatalf("unexpected error: %s", err)
+		}
+		query := fmt.Sprintf("SELECT_client_close_after_%d", time.Now().UnixNano())
+
+		closeA := make(chan bool, 1)
+		doneA := make(chan struct{})
+		go func() {
+			defer close(doneA)
+			proxy.ServeHTTP(&closeNotifierOnSignal{ResponseWriter: httptest.NewRecorder(), closed: closeA}, newRequest(query))
+		}()
+		waitFor(t, "the first query to run on the server", func() bool {
+			return atomic.LoadInt64(&nbRequestsInflight) == 1
+		})
+		closeA <- true
+		<-doneA
+
+		// The transaction state lives for a short time after the end of A: the same query sent
+		// at once must run, not get the error of the cancelled one.
+		assertQueryRan(t, makeCustomRequest(proxy, newRequest(query)))
+	})
+}
+
 func TestKillQuery(t *testing.T) {
 	testCases := []struct {
 		name string
diff --git a/scope.go b/scope.go
index ef970565..f1caef37 100644
--- a/scope.go
+++ b/scope.go
@@ -52,6 +52,11 @@ type scope struct {
 	// is true when KillQuery has been called
 	canceled bool
 
+	// is true when proxyRequest ended with context.Canceled. In that path only listenToCloseNotify
+	// cancels the context, so it means that the client closed the connection; timeouts end with
+	// context.DeadlineExceeded.
+	clientClosed bool
+
 	labels prometheus.Labels
 
 	requestPacketSize int