This is an automated email from the ASF dual-hosted git repository.
joao-r-reis pushed a commit to branch trunk
in repository https://gitbox.apache.org/repos/asf/cassandra-gocql-driver.git
The following commit(s) were added to refs/heads/trunk by this push:
new 60ef3281 Add request attempt interceptors for queries and batches
60ef3281 is described below
commit 60ef32811be61529676d2ebe64ba11a856443e83
Author: Alex Benn <[email protected]>
AuthorDate: Tue Oct 20 16:37:06 2026 -0400
Add request attempt interceptors for queries and batches
Applications need to apply rate limiting, propagate tracing context, and
inspect results for each request attempt, including retries and speculative
executions. Observers report results after an attempt and cannot reject
execution or adjust its context.
Introduce pre- and post-attempt hooks with cluster defaults and per-query
or per-batch overrides. Interceptor chains preserve each interceptor's
context and invoke post-attempt hooks in reverse order.
Track attempt indices separately from completed attempts so observers
receive stable indices during concurrent execution while retry decisions
and average latency retain their existing semantics.
Patch by Ben Eddy, Alex Benn, João Reis; reviewed by João Reis, Bohdan
Siryk, toptobes for CASSGO-39
Assisted-by: Claude Code:claude-opus-5-5
---
CHANGELOG.md | 6 ++
cluster.go | 4 +
control.go | 6 +-
doc.go | 12 +++
example_interceptor_test.go | 154 ++++++++++++++++++++++++++++
policies_test.go | 4 +-
query_executor.go | 237 ++++++++++++++++++++++++++++++++++++++------
query_executor_unit_test.go | 178 +++++++++++++++++++++++++++++++++
session.go | 75 ++++++++++++--
9 files changed, 632 insertions(+), 44 deletions(-)
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 071f2e93..8493442d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,12 @@ All notable changes to this project will be documented in this
file.
The format is based on [Keep a
Changelog](https://keepachangelog.com/en/1.0.0/),
and this project adheres to [Semantic
Versioning](https://semver.org/spec/v2.0.0.html).
+## [Unreleased]
+
+### Added
+
+- Add request attempt interceptors for queries and batches (CASSGO-39)
+
## [2.2.0]
### Added
diff --git a/cluster.go b/cluster.go
index eeb12b92..f2f8dd1b 100644
--- a/cluster.go
+++ b/cluster.go
@@ -258,6 +258,10 @@ type ClusterConfig struct {
// See https://issues.apache.org/jira/browse/CASSANDRA-10786
DisableSkipMetadata bool
+ // RequestInterceptor will set the provided interceptor on all
queries/batches created from this session.
+ // Use it to intercept requests by providing an implementation of
RequestInterceptor.
+ RequestInterceptor RequestInterceptor
+
// QueryObserver will set the provided query observer on all queries
created from this session.
// Use it to collect metrics / stats from queries by providing an
implementation of QueryObserver.
QueryObserver QueryObserver
diff --git a/control.go b/control.go
index 9f56958a..5acca9ba 100644
--- a/control.go
+++ b/control.go
@@ -591,7 +591,11 @@ func (c *controlConn) query(statement string, values
...interface{}) (iter *Iter
NewLogFieldString("statement", statement),
NewLogFieldError("err", iter.err))
}
- qry.metrics.attempt(0)
+ // Record each attempt so the retry policy sees an increasing
attempt count and
+ // eventually stops retrying. getNextAttempt() is called to
keep attemptsStarted
+ // in sync with attemptsCompleted, as queryMetrics assumes
started >= completed.
+ qry.metrics.getNextAttempt()
+ qry.metrics.recordAttempt(0)
qry.hostMetricsManager.attempt(0, c.getConn().host)
if iter.err == nil || !c.retry.Attempt(qry) {
break
diff --git a/doc.go b/doc.go
index c0b9a729..a05a5e33 100644
--- a/doc.go
+++ b/doc.go
@@ -813,6 +813,18 @@
//
// See Example_userDefinedTypesMap, Example_userDefinedTypesStruct,
ExampleUDTMarshaler, ExampleUDTUnmarshaler.
//
+// # Interceptors
+//
+// A RequestInterceptor wraps query/batch execution and can be used to inject
logic that should apply to all request
+// execution attempts. For example, interceptors can be used for rate
limiting, logging, attaching distributed tracing
+// metadata to the context, and inspecting request results. However, the
query/batch itself cannot be modified.
+//
+// A RequestInterceptor will be invoked once prior to (via
InterceptPreAttempt()) and once after (via
+// InterceptPostAttempt()) each request execution attempt, including retry
attempts and speculative execution attempts.
+// It may optionally return an error to prevent that request from being
executed.
+//
+// See Example_interceptorFault, Example_interceptorContext, and
Example_interceptorChain for examples.
+//
// # Metrics and tracing
//
// It is possible to provide observer implementations that could be used to
gather metrics:
diff --git a/example_interceptor_test.go b/example_interceptor_test.go
new file mode 100644
index 00000000..370d7bfc
--- /dev/null
+++ b/example_interceptor_test.go
@@ -0,0 +1,154 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+/*
+ * Content before git sha 34fdeebefcbf183ed7f916f931aa0586fdaa1b40
+ * Copyright (c) 2016, The Gocql authors,
+ * provided under the BSD-3-Clause License.
+ * See the NOTICE file distributed with this work for additional information.
+ */
+
+package gocql_test
+
+import (
+ "context"
+ "errors"
+ "log"
+
+ gocql "github.com/apache/cassandra-gocql-driver/v2"
+)
+
+type MyRequestInterceptor struct {
+ injectFault bool
+ contextStr string
+}
+
+type traceIDKey struct{}
+
+var errInjectedTimeout = errors.New("timeout")
+
+var _ gocql.RequestInterceptor = (*MyRequestInterceptor)(nil)
+
+func (q MyRequestInterceptor) InterceptPreAttempt(
+ ctx context.Context,
+ attempt gocql.RequestAttempt,
+) (gocql.InterceptResult, error) {
+ switch attempt.Type {
+ case gocql.StatementQuery:
+ // Inspect query
+ log.Println(attempt.Query.Statement())
+ case gocql.StatementBatch:
+ // Inspect batch
+ log.Println(attempt.Batch.Entries[0].Stmt)
+ }
+
+ // Inspect or modify context
+ ctx = context.WithValue(ctx, traceIDKey{}, q.contextStr)
+
+ // Optionally bypass request issuance by returning an error to prevent
query execution.
+ // For example, to simulate query timeouts, or to perform client-side
load shedding.
+ if q.injectFault {
+ return gocql.InterceptResult{}, errInjectedTimeout
+ }
+
+ return gocql.InterceptResult{Ctx: ctx}, nil
+}
+
+func (q MyRequestInterceptor) InterceptPostAttempt(res gocql.InterceptResult,
_ error) {
+ if res.Ctx == nil {
+ return
+ }
+ if q.contextStr != res.Ctx.Value(traceIDKey{}) {
+ log.Printf("context string changed: %s",
res.Ctx.Value(traceIDKey{}))
+ }
+}
+
+// Demonstrates how a RequestInterceptor might be used to inject faults.
+func Example_interceptorFault() {
+ cluster := gocql.NewCluster("localhost:9042")
+ cluster.RequestInterceptor = MyRequestInterceptor{injectFault: true}
+
+ session, err := cluster.CreateSession()
+ if err != nil {
+ log.Fatal(err)
+ }
+ defer session.Close()
+
+ ctx := context.Background()
+
+ var stringValue string
+ // Here we specify a retry policy, but the GoCQL library with
SimpleRetryPolicy will only retry multiple
+ // times if you have more than one host. This may not be the case if
you are running this example against
+ // a local test Cassandra instance, so NumRetries: 2 is here only for
explanatory purposes.
+ err = session.Query("select now() from system.local").
+ RetryPolicy(&gocql.SimpleRetryPolicy{NumRetries: 2}).
+ Idempotent(true).
+ ScanContext(ctx, &stringValue)
+ if !errors.Is(err, errInjectedTimeout) {
+ log.Fatal(err)
+ }
+}
+
+// Demonstrates how a RequestInterceptor might be used to propagate ancillary
data to post-processing via the context.
+func Example_interceptorContext() {
+ cluster := gocql.NewCluster("localhost:9042")
+ cluster.RequestInterceptor = MyRequestInterceptor{injectFault: false}
+
+ session, err := cluster.CreateSession()
+ if err != nil {
+ log.Fatal(err)
+ }
+ defer session.Close()
+
+ ctx := context.Background()
+
+ var stringValue string
+ err = session.Query("select now() from system.local").
+ RetryPolicy(&gocql.SimpleRetryPolicy{NumRetries: 2}).
+ ScanContext(ctx, &stringValue)
+ if err != nil {
+ log.Fatal(err)
+ }
+}
+
+// Demonstrates how to use a request interceptor chain.
+func Example_interceptorChain() {
+ cluster := gocql.NewCluster("localhost:9042")
+ cluster.RequestInterceptor = gocql.RequestInterceptorChain{
+ []gocql.RequestInterceptor{
+ MyRequestInterceptor{contextStr: "123"},
+ MyRequestInterceptor{contextStr: "234"},
+ MyRequestInterceptor{contextStr: "345"},
+ },
+ }
+
+ session, err := cluster.CreateSession()
+ if err != nil {
+ log.Fatal(err)
+ }
+ defer session.Close()
+
+ ctx := context.Background()
+
+ var stringValue string
+ err = session.Query("select now() from system.local").
+ RetryPolicy(&gocql.SimpleRetryPolicy{NumRetries: 2}).
+ ScanContext(ctx, &stringValue)
+ if err != nil {
+ log.Fatal(err)
+ }
+}
diff --git a/policies_test.go b/policies_test.go
index 547a123b..aaae1925 100644
--- a/policies_test.go
+++ b/policies_test.go
@@ -254,7 +254,7 @@ func TestSimpleRetryPolicy(t *testing.T) {
}
for _, c := range cases {
- q.metrics = &queryMetrics{totalAttempts: int64(c.attempts)}
+ q.metrics = &queryMetrics{attemptsCompleted: c.attempts}
q.hostMetricsManager =
preFilledHostMetricsMetricsManager(map[string]*hostMetrics{"127.0.0.1":
{Attempts: c.attempts}})
if c.allow && !rt.Attempt(q) {
t.Fatalf("should allow retry after %d attempts",
c.attempts)
@@ -339,7 +339,7 @@ func TestDowngradingConsistencyRetryPolicy(t *testing.T) {
}
for _, c := range cases {
- q.metrics = &queryMetrics{totalAttempts: int64(c.attempts)}
+ q.metrics = &queryMetrics{attemptsCompleted: c.attempts}
q.hostMetricsManager =
preFilledHostMetricsMetricsManager(map[string]*hostMetrics{"127.0.0.1":
{Attempts: c.attempts}})
if c.retryType != rt.GetRetryType(c.err) {
t.Fatalf("retry type should be %v", c.retryType)
diff --git a/query_executor.go b/query_executor.go
index 5d848c2e..1ab75a4c 100644
--- a/query_executor.go
+++ b/query_executor.go
@@ -58,7 +58,8 @@ type Statement interface {
type internalRequest interface {
execute(ctx context.Context, conn *Conn) *Iter
- attempt(keyspace string, end, start time.Time, iter *Iter, host
*HostInfo)
+ getNextAttempt() int
+ recordAttempt(attemptNum int, keyspace string, end, start time.Time,
iter *Iter, host *HostInfo)
retryPolicy() RetryPolicy
speculativeExecutionPolicy() SpeculativeExecutionPolicy
getQueryMetrics() *queryMetrics
@@ -73,12 +74,164 @@ type queryExecutor struct {
policy HostSelectionPolicy
}
-func (q *queryExecutor) attemptQuery(ctx context.Context, qry internalRequest,
conn *Conn) *Iter {
+type StatementType int
+
+const (
+ StatementQuery StatementType = iota
+ StatementBatch
+)
+
+type RequestAttempt struct {
+ // The statement type, whether query or batch.
+ Type StatementType
+ // Only one of Query or Batch will be set, depending on the op type.
Both should be considered
+ // immutable.
+ Query *Query
+ Batch *Batch
+ // The host that will receive the query.
+ Host *HostInfo
+}
+
+type InterceptResult struct {
+ Ctx context.Context
+}
+
+// RequestInterceptor is the interface implemented by interceptors /
middleware.
+//
+// Interceptors are well-suited to logic that is not specific to a single
query or batch, such as flow control.
+type RequestInterceptor interface {
+ // InterceptPreAttempt is invoked once immediately before a query or
batch execution attempt, including retry
+ // attempts and speculative execution attempts.
+ //
+ // The method may inspect the RequestAttempt and its contents (e.g.
Query/Batch, connection information, etc.), along
+ // with the context ctx. If a ctx is specified in the returned
InterceptResult, it is used when performing the
+ // execution, offering a way to adjust deadline or other metadata. To
prevent request execution, return a non-nil
+ // error.
+ InterceptPreAttempt(ctx context.Context, attempt RequestAttempt)
(InterceptResult, error)
+
+ // InterceptPostAttempt is called with the intercept result created at
pre-attempt time, along with any
+ // request execution error. Note that if this attempt failed, another
attempt in flight could still succeed.
+ // Must not panic.
+ // N.B.: InterceptPostAttempt is not called if InterceptPreAttempt
returns an error.
+ InterceptPostAttempt(res InterceptResult, execResult error)
+}
+
+// RequestInterceptorChain can be used to apply multiple Interceptors at
Intercept time.
+// Each interceptor is passed the previous one's context if that previous one
returned a non-nil
+// context, and the final request attempt is performed with the last
interceptor's context.
+//
+// InterceptPreAttempt is called on the Interceptors in order, and
InterceptPostAttempt is called
+// in reverse order. At post-attempt time, Interceptors are each handed back
the context they synthesized at
+// pre-attempt time, rather than the final context. This is to avoid confusion
if e.g. a pre-attempt grabbed some
+// resource that must be released, and multiple interceptors might use the
same context key to attach that
+// resource to the attempt's request. However, this also means that if an
Interceptor returns an InterceptResult
+// with a nil context, they will be handed back a nil.
+//
+// Finally, if Interceptor i's InterceptPreAttempt returns an error,
indicating that execution should not occur,
+// InterceptPostAttempt is still called on all the interceptors that were
chained before i. They are given the error
+// that the failing InterceptPreAttempt returned, in order to signal that
execution was not successful.
+type RequestInterceptorChain struct {
+ Interceptors []RequestInterceptor
+}
+
+type interceptorChainContextList []context.Context
+
+type interceptorChainContextListKey struct{}
+
+var _ RequestInterceptor = (*RequestInterceptorChain)(nil)
+
+func (c RequestInterceptorChain) InterceptPreAttempt(
+ ctx context.Context,
+ attempt RequestAttempt,
+) (InterceptResult, error) {
+ capturedContexts := make([]context.Context, 0, len(c.Interceptors))
+ lastCtx := ctx
+ for i, interceptor := range c.Interceptors {
+ result, err := interceptor.InterceptPreAttempt(lastCtx, attempt)
+ if err != nil {
+ // We must unwind all previously-successful
interceptors before returning err, since
+ // those interceptors may be counting on
InterceptPostAttempt being invoked.
+ for j := i - 1; j >= 0; j-- {
+ // Pass in the final err to all PostAttempt
intercepts, to signal that execution was not successful
+ // (because it never actually happened).
+
c.Interceptors[j].InterceptPostAttempt(InterceptResult{Ctx:
capturedContexts[j]}, err)
+ }
+ return InterceptResult{}, err
+ }
+ capturedContexts = append(capturedContexts, result.Ctx)
+ if result.Ctx != nil {
+ lastCtx = result.Ctx
+ }
+ }
+
+ finalContext := context.WithValue(lastCtx,
interceptorChainContextListKey{}, interceptorChainContextList(capturedContexts))
+ // Returns a result whose context wraps all the captured contexts for
all the chained interceptors.
+ return InterceptResult{Ctx: finalContext}, nil
+}
+
+func (c RequestInterceptorChain) InterceptPostAttempt(
+ res InterceptResult, execResult error,
+) {
+ // res must be the result returned by this chain's InterceptPreAttempt,
which carries each interceptor's context.
+ // If it isn't, those contexts are lost, so skip the post-attempt calls
rather than panic.
+ if res.Ctx == nil {
+ return
+ }
+ ctxList, ok :=
res.Ctx.Value(interceptorChainContextListKey{}).(interceptorChainContextList)
+ if !ok {
+ return
+ }
+ for i := len(ctxList) - 1; i >= 0; i-- {
+ interceptor := c.Interceptors[i]
+ ctx := ctxList[i]
+ interceptor.InterceptPostAttempt(InterceptResult{Ctx: ctx},
execResult)
+ }
+}
+
+func (q *queryExecutor) attemptQuery(ctx context.Context, iRequest
internalRequest, conn *Conn) *Iter {
+ var res InterceptResult
+ var interceptor RequestInterceptor
+ iq := RequestAttempt{
+ Host: conn.host,
+ }
+ if iQuery, ok := iRequest.(*internalQuery); ok {
+ iq.Type = StatementQuery
+ iq.Query = iQuery.originalQuery
+ interceptor = iQuery.qryOpts.interceptor
+ }
+ if iBatch, ok := iRequest.(*internalBatch); ok {
+ iq.Type = StatementBatch
+ iq.Batch = iBatch.originalBatch
+ interceptor = iBatch.batchOpts.interceptor
+ }
+ attempt := iRequest.getNextAttempt()
+ if interceptor != nil {
+ var err error
+ res, err = interceptor.InterceptPreAttempt(ctx, iq)
+ if err != nil {
+ iter := newErrIter(err, iRequest.getQueryMetrics(),
iRequest.Keyspace(), iRequest.getRoutingInfo(), iRequest.getKeyspaceFunc())
+ iter.rejectedByInterceptor = true
+ // Always record the attempt, even if the interceptor
aborts it before it is actually executed. This way, a
+ // custom retry policy will see the retry count
increase and respond accordingly if the interceptor decides
+ // to abort the attempt, rather than getting into an
infinite retry loop with zero perceived attempts.
+ // Nothing was sent, so the attempt has zero latency.
+ now := time.Now()
+ iRequest.recordAttempt(attempt, q.pool.keyspace, now,
now, iter, conn.host)
+ return iter
+ }
+ if res.Ctx != nil {
+ ctx = res.Ctx
+ }
+ }
+
+ // Time only the execution, so that time spent in interceptors isn't
reported as query latency.
start := time.Now()
- iter := qry.execute(ctx, conn)
+ iter := iRequest.execute(ctx, conn)
end := time.Now()
-
- qry.attempt(q.pool.keyspace, end, start, iter, conn.host)
+ if interceptor != nil {
+ interceptor.InterceptPostAttempt(res, iter.err)
+ }
+ iRequest.recordAttempt(attempt, q.pool.keyspace, end, start, iter,
conn.host)
return iter
}
@@ -199,10 +352,16 @@ func (q *queryExecutor) do(ctx context.Context, qry
internalRequest, hostIter Ne
case context.Canceled, context.DeadlineExceeded, ErrNotFound:
// those errors represents logical errors, they should
not count
// toward removing a node from the pool
- selectedHost.Mark(nil)
+ if !iter.rejectedByInterceptor {
+ // However, if the error comes from an
interceptor, it's not
+ // actually a success, so we don't Mark(nil) in
that case.
+ selectedHost.Mark(nil)
+ }
return iter
default:
- selectedHost.Mark(iter.err)
+ if !iter.rejectedByInterceptor {
+ selectedHost.Mark(iter.err)
+ }
}
// Exit if the query was successful
@@ -262,9 +421,10 @@ type queryOptions struct {
pageSize int
disableAutoPage bool
- // Monitoring
- trace Tracer
- observer QueryObserver
+ // Monitoring and intercepting
+ trace Tracer
+ observer QueryObserver
+ interceptor RequestInterceptor
// Parameters
values []interface{}
@@ -313,6 +473,7 @@ func newQueryOptions(q *Query, ctx context.Context)
*queryOptions {
prefetch: q.prefetch,
trace: q.trace,
observer: q.observer,
+ interceptor: q.interceptor,
rt: q.rt,
spec: q.spec,
binding: q.binding,
@@ -371,14 +532,18 @@ func newInternalQuery(q *Query, ctx context.Context)
*internalQuery {
}
}
-// Attempts returns the number of times the query was executed.
+// getNextAttempt returns the index of the next attempt. Calling this
increments the attempt counter by one.
+func (q *internalQuery) getNextAttempt() int {
+ return q.metrics.getNextAttempt()
+}
+
func (q *internalQuery) Attempts() int {
return q.metrics.attempts()
}
-func (q *internalQuery) attempt(keyspace string, end, start time.Time, iter
*Iter, host *HostInfo) {
+func (q *internalQuery) recordAttempt(attemptNum int, keyspace string, end,
start time.Time, iter *Iter, host *HostInfo) {
latency := end.Sub(start)
- attempt := q.metrics.attempt(latency)
+ q.metrics.recordAttempt(latency)
if q.qryOpts.observer != nil {
metricsForHost := q.hostMetricsManager.attempt(latency, host)
@@ -389,16 +554,17 @@ func (q *internalQuery) attempt(keyspace string, end,
start time.Time, iter *Ite
Keyspace: q.routingInfo.getKeyspace(),
Table: q.routingInfo.getTable(),
},
- IsPrepared: q.routingInfo.isPrepared(),
- Values: q.qryOpts.values,
- Start: start,
- End: end,
- Rows: iter.numRows,
- Host: host,
- Metrics: metricsForHost,
- Err: iter.err,
- Attempt: attempt,
- Query: q.originalQuery,
+ IsPrepared: q.routingInfo.isPrepared(),
+ Values: q.qryOpts.values,
+ Start: start,
+ End: end,
+ Rows: iter.numRows,
+ Host: host,
+ Metrics: metricsForHost,
+ Err: iter.err,
+ Attempt: attemptNum,
+ Query: q.originalQuery,
+ RejectedByInterceptor: iter.rejectedByInterceptor,
})
}
}
@@ -505,8 +671,9 @@ func (q *internalQuery) getKeyspaceFunc() func() string {
}
type batchOptions struct {
- trace Tracer
- observer BatchObserver
+ trace Tracer
+ observer BatchObserver
+ interceptor RequestInterceptor
bType BatchType
entries []BatchEntry
@@ -549,6 +716,7 @@ func newBatchOptions(b *Batch, ctx context.Context)
*batchOptions {
spec: b.spec,
trace: b.trace,
observer: b.observer,
+ interceptor: b.interceptor,
serialCons: b.serialCons,
defaultTimestamp: b.defaultTimestamp,
defaultTimestampValue: b.defaultTimestampValue,
@@ -601,14 +769,18 @@ func newInternalBatch(batch *Batch, ctx context.Context)
*internalBatch {
}
}
+func (b *internalBatch) getNextAttempt() int {
+ return b.metrics.getNextAttempt()
+}
+
// Attempts returns the number of attempts made to execute the batch.
func (b *internalBatch) Attempts() int {
return b.metrics.attempts()
}
-func (b *internalBatch) attempt(keyspace string, end, start time.Time, iter
*Iter, host *HostInfo) {
+func (b *internalBatch) recordAttempt(attemptNum int, keyspace string, end,
start time.Time, iter *Iter, host *HostInfo) {
latency := end.Sub(start)
- attempt := b.metrics.attempt(latency)
+ b.metrics.recordAttempt(latency)
if b.batchOpts.observer == nil {
return
@@ -641,11 +813,12 @@ func (b *internalBatch) attempt(keyspace string, end,
start time.Time, iter *Ite
Start: start,
End: end,
// Rows not used in batch observations // TODO - might be able
to support it when using BatchCAS
- Host: host,
- Metrics: metricsForHost,
- Err: iter.err,
- Attempt: attempt,
- Batch: b.originalBatch,
+ Host: host,
+ Metrics: metricsForHost,
+ Err: iter.err,
+ Attempt: attemptNum,
+ Batch: b.originalBatch,
+ RejectedByInterceptor: iter.rejectedByInterceptor,
})
}
diff --git a/query_executor_unit_test.go b/query_executor_unit_test.go
new file mode 100644
index 00000000..b6693b65
--- /dev/null
+++ b/query_executor_unit_test.go
@@ -0,0 +1,178 @@
+//go:build all || unit
+// +build all unit
+
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package gocql
+
+import (
+ "context"
+ "errors"
+ "sync/atomic"
+ "testing"
+ "time"
+)
+
+// markCountingPolicy wraps a HostSelectionPolicy and counts Mark calls on the
hosts it picks.
+type markCountingPolicy struct {
+ HostSelectionPolicy
+ marks int32
+}
+
+func (p *markCountingPolicy) Pick(qry ExecutableStatement) NextHost {
+ next := p.HostSelectionPolicy.Pick(qry)
+ return func() SelectedHost {
+ host := next()
+ if host == nil {
+ return nil
+ }
+ return markCountingHost{SelectedHost: host, marks: &p.marks}
+ }
+}
+
+type markCountingHost struct {
+ SelectedHost
+ marks *int32
+}
+
+func (h markCountingHost) Mark(err error) {
+ atomic.AddInt32(h.marks, 1)
+ h.SelectedHost.Mark(err)
+}
+
+type rejectingInterceptor struct {
+ err error
+}
+
+func (r rejectingInterceptor) InterceptPreAttempt(context.Context,
RequestAttempt) (InterceptResult, error) {
+ return InterceptResult{}, r.err
+}
+
+func (rejectingInterceptor) InterceptPostAttempt(InterceptResult, error) {}
+
+// timingInterceptor records when its pre-attempt hook returns and its
post-attempt hook is called.
+type timingInterceptor struct {
+ preDone time.Time
+ postStart time.Time
+}
+
+func (ti *timingInterceptor) InterceptPreAttempt(ctx context.Context, _
RequestAttempt) (InterceptResult, error) {
+ ti.preDone = time.Now()
+ return InterceptResult{Ctx: ctx}, nil
+}
+
+func (ti *timingInterceptor) InterceptPostAttempt(InterceptResult, error) {
+ ti.postStart = time.Now()
+}
+
+type countingInterceptor struct {
+ posts int
+}
+
+func (c *countingInterceptor) InterceptPreAttempt(ctx context.Context, _
RequestAttempt) (InterceptResult, error) {
+ return InterceptResult{Ctx: ctx}, nil
+}
+
+func (c *countingInterceptor) InterceptPostAttempt(InterceptResult, error) {
+ c.posts++
+}
+
+type lastQueryObserver struct {
+ query ObservedQuery
+}
+
+func (o *lastQueryObserver) ObserveQuery(_ context.Context, q ObservedQuery) {
+ o.query = q
+}
+
+func TestInterceptorRejectionDoesNotMarkHost(t *testing.T) {
+ srv := NewTestServer(t, defaultProto, context.Background())
+ defer srv.Stop()
+
+ policy := &markCountingPolicy{HostSelectionPolicy:
RoundRobinHostPolicy()}
+ cluster := testCluster(defaultProto, srv.Address)
+ cluster.PoolConfig.HostSelectionPolicy = policy
+ db, err := cluster.CreateSession()
+ if err != nil {
+ t.Fatalf("NewCluster: %v", err)
+ }
+ defer db.Close()
+
+ // A sent attempt marks the host, which shows the policy is in use.
+ if err := db.Query("void").Exec(); err != nil {
+ t.Fatal(err)
+ }
+ if marks := atomic.LoadInt32(&policy.marks); marks != 1 {
+ t.Fatalf("expected 1 mark after a sent attempt, got %d", marks)
+ }
+
+ errRejected := errors.New("rejected")
+ iter := db.Query("void").Interceptor(rejectingInterceptor{err:
errRejected}).Iter()
+ if err := iter.Close(); !errors.Is(err, errRejected) {
+ t.Fatalf("expected %v, got %v", errRejected, err)
+ }
+ if marks := atomic.LoadInt32(&policy.marks); marks != 1 {
+ t.Fatalf("expected a rejected attempt not to mark the host, got
%d marks", marks)
+ }
+}
+
+func TestInterceptorTimeExcludedFromLatency(t *testing.T) {
+ srv := NewTestServer(t, defaultProto, context.Background())
+ defer srv.Stop()
+
+ db, err := newTestSession(defaultProto, srv.Address)
+ if err != nil {
+ t.Fatalf("NewCluster: %v", err)
+ }
+ defer db.Close()
+
+ interceptor := &timingInterceptor{}
+ observer := &lastQueryObserver{}
+ if err :=
db.Query("void").Interceptor(interceptor).Observer(observer).Exec(); err != nil
{
+ t.Fatal(err)
+ }
+
+ if observer.query.Start.Before(interceptor.preDone) {
+ t.Errorf("attempt timing started before the pre-attempt
interceptor returned")
+ }
+ if observer.query.End.After(interceptor.postStart) {
+ t.Errorf("attempt timing ended after the post-attempt
interceptor was called")
+ }
+}
+
+func TestInterceptorChainPostAttemptWithoutChainContext(t *testing.T) {
+ inner := &countingInterceptor{}
+ chain := RequestInterceptorChain{Interceptors:
[]RequestInterceptor{inner}}
+
+ // Neither result came from chain.InterceptPreAttempt, so there are no
contexts to hand back.
+ chain.InterceptPostAttempt(InterceptResult{}, nil)
+ chain.InterceptPostAttempt(InterceptResult{Ctx: context.Background()},
nil)
+ if inner.posts != 0 {
+ t.Fatalf("expected no post-attempt calls without the chain's
context, got %d", inner.posts)
+ }
+
+ res, err := chain.InterceptPreAttempt(context.Background(),
RequestAttempt{})
+ if err != nil {
+ t.Fatal(err)
+ }
+ chain.InterceptPostAttempt(res, nil)
+ if inner.posts != 1 {
+ t.Fatalf("expected 1 post-attempt call, got %d", inner.posts)
+ }
+}
diff --git a/session.go b/session.go
index bfe42573..feb0752a 100644
--- a/session.go
+++ b/session.go
@@ -62,6 +62,7 @@ type Session struct {
connectObserver ConnectObserver
frameObserver FrameHeaderObserver
streamObserver StreamObserver
+ interceptor RequestInterceptor
hostSource *ringDescriber
ringRefresher *refreshDebouncer
stmtsLRU *preparedLRU
@@ -194,6 +195,7 @@ func NewSession(cfg ClusterConfig) (*Session, error) {
s.connectObserver = cfg.ConnectObserver
s.frameObserver = cfg.FrameHeaderObserver
s.streamObserver = cfg.StreamObserver
+ s.interceptor = cfg.RequestInterceptor
// Propagate node status, topology and schema change listeners
s.hostListeners = newInternalHostStateListeners(
@@ -973,26 +975,50 @@ type hostMetrics struct {
TotalLatency int64
}
+// queryMetrics tracks attempt counts and latency for a query or batch.
+//
+// When speculative execution is active, attempts can start and finish out of
order
+// across concurrent goroutines. Attempt indices are issued at start time via
getNextAttempt()
+// so observers receive stable attempt numbers, while attemptsCompleted and
totalLatency
+// are recorded at completion time under lock so latency() (average latency)
and attempts()
+// remain consistent mid-flight.
type queryMetrics struct {
- totalAttempts int64
- totalLatency int64
+ l sync.RWMutex
+ attemptsStarted int
+ attemptsCompleted int
+ totalLatency int64
}
-func (qm *queryMetrics) attempt(addLatency time.Duration) int {
- atomic.AddInt64(&qm.totalLatency, addLatency.Nanoseconds())
- return int(atomic.AddInt64(&qm.totalAttempts, 1) - 1)
+func (qm *queryMetrics) getNextAttempt() int {
+ qm.l.Lock()
+ defer qm.l.Unlock()
+ attempt := qm.attemptsStarted
+ qm.attemptsStarted++
+ return attempt
}
+func (qm *queryMetrics) recordAttempt(addLatency time.Duration) {
+ qm.l.Lock()
+ defer qm.l.Unlock()
+ qm.totalLatency += addLatency.Nanoseconds()
+ qm.attemptsCompleted++
+}
+
+// Returns total number of attempts completed.
func (qm *queryMetrics) attempts() int {
- return int(atomic.LoadInt64(&qm.totalAttempts))
+ qm.l.RLock()
+ defer qm.l.RUnlock()
+ return qm.attemptsCompleted
}
func (qm *queryMetrics) latency() int64 {
- attempts := atomic.LoadInt64(&qm.totalAttempts)
+ qm.l.RLock()
+ defer qm.l.RUnlock()
+ attempts := qm.attemptsCompleted
if attempts == 0 {
- return atomic.LoadInt64(&qm.totalLatency)
+ return qm.totalLatency
}
- return atomic.LoadInt64(&qm.totalLatency) / attempts
+ return qm.totalLatency / int64(attempts)
}
type hostMetricsManager interface {
@@ -1055,6 +1081,7 @@ type Query struct {
prefetch float64
trace Tracer
observer QueryObserver
+ interceptor RequestInterceptor
session *Session
rt RetryPolicy
spec SpeculativeExecutionPolicy
@@ -1120,6 +1147,7 @@ func (q *Query) defaultsFromSession() {
q.pageSize = s.pageSize
q.trace = s.trace
q.observer = s.queryObserver
+ q.interceptor = s.interceptor
q.prefetch = s.prefetch
q.rt = s.cfg.RetryPolicy
q.serialCons = s.cfg.SerialConsistency
@@ -1194,6 +1222,14 @@ func (q *Query) Observer(observer QueryObserver) *Query {
return q
}
+// Interceptor specifies a query-level interceptor on this query.
+// The provided interceptor will be called every time this query is executed.
+// This replaces any session-level inherited interceptor for this query.
+func (q *Query) Interceptor(interceptor RequestInterceptor) *Query {
+ q.interceptor = interceptor
+ return q
+}
+
// PageSize will tell the iterator to fetch the result in pages of size n.
// This is useful for iterating over large result sets, but setting the
// page size too low might decrease the performance. This feature is only
@@ -1602,6 +1638,9 @@ type Iter struct {
host *HostInfo
metrics *queryMetrics
+ // rejectedByInterceptor is set when a RequestInterceptor aborted the
attempt before it was sent.
+ rejectedByInterceptor bool
+
getKeyspace func() string
keyspace string
routingInfo *queryRoutingInfo
@@ -2013,6 +2052,7 @@ type Batch struct {
spec SpeculativeExecutionPolicy
trace Tracer
observer BatchObserver
+ interceptor RequestInterceptor
session *Session
serialCons Consistency
defaultTimestamp bool
@@ -2038,6 +2078,7 @@ func (s *Session) Batch(typ BatchType) *Batch {
serialCons: s.cfg.SerialConsistency,
trace: s.trace,
observer: s.batchObserver,
+ interceptor: s.interceptor,
session: s,
Cons: s.cons,
defaultTimestamp: s.cfg.DefaultTimestamp,
@@ -2062,6 +2103,14 @@ func (b *Batch) Observer(observer BatchObserver) *Batch {
return b
}
+// Interceptor specifies a batch-level interceptor on this batch.
+// The provided interceptor will be called every time this batched query is
executed.
+// This replaces any session-level interceptor for this batch.
+func (b *Batch) Interceptor(interceptor RequestInterceptor) *Batch {
+ b.interceptor = interceptor
+ return b
+}
+
func (b *Batch) Keyspace() string {
return b.keyspace
}
@@ -2441,6 +2490,10 @@ type ObservedQuery struct {
// Query object associated with this request. Should be used as read
only.
Query *Query
+
+ // Whether the attempt was aborted by the interceptor before it was
issued. In this case, the request was never
+ // actually sent to the above host.
+ RejectedByInterceptor bool
}
// QueryObserver is the interface implemented by query observers / stat
collectors.
@@ -2490,6 +2543,10 @@ type ObservedBatch struct {
// Batch object associated with this request. Should be used as read
only.
Batch *Batch
+
+ // Whether the attempt was aborted by the interceptor before it was
issued. In this case, the request was never
+ // actually sent to the above host.
+ RejectedByInterceptor bool
}
// BatchObserver is the interface implemented by batch observers / stat
collectors.
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
