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