diff --git a/cluster.go b/cluster.go index 0c687d20d..5c0bcbe9d 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 + // ExecAttemptInterceptor will set the provided interceptor on all queries/batches created from this session. + // Use it to intercept queries by providing an implementation of ExecAttemptInterceptor. + ExecAttemptInterceptor ExecAttemptInterceptor + // 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 9f56958a3..b1dda91f1 100644 --- a/control.go +++ b/control.go @@ -591,7 +591,7 @@ func (c *controlConn) query(statement string, values ...interface{}) (iter *Iter NewLogFieldString("statement", statement), NewLogFieldError("err", iter.err)) } - qry.metrics.attempt(0) + qry.metrics.recordAttempt(0, 0, nil) 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 c0b9a7299..7035d2b8a 100644 --- a/doc.go +++ b/doc.go @@ -813,6 +813,19 @@ // // See Example_userDefinedTypesMap, Example_userDefinedTypesStruct, ExampleUDTMarshaler, ExampleUDTUnmarshaler. // +// # Interceptors +// +// A ExecAttemptInterceptor wraps query/batch execution and can be used to inject logic that should apply to all query +// and batch execution attempts. For example, interceptors can be used for rate limiting, logging, attaching +// distributed tracing metadata to the context, and inspecting query/batch results. However, the query/batch itself +// cannot be modified. +// +// A ExecAttemptInterceptor will be invoked once prior to each query execution attempt, including retry attempts +// and speculative execution attempts. Interceptors are responsible for calling the provided handler and returning +// a non-nil Iter or an error. +// +// See Example_interceptor for full example. +// // # 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 000000000..eda93ab09 --- /dev/null +++ b/example_interceptor_test.go @@ -0,0 +1,115 @@ +/* + * 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" + "log" + "time" + + gocql "github.com/apache/cassandra-gocql-driver/v2" +) + +type MyExecAttemptInterceptor struct { + injectFault bool +} + +func (q MyExecAttemptInterceptor) Intercept( + ctx context.Context, + attempt gocql.QueryAttempt, + handler gocql.QueryAttemptHandler, +) (*gocql.Iter, error) { + switch attempt.Type { + case gocql.OpQuery: + // Inspect query + log.Println(attempt.Query.Statement()) + case gocql.OpBatch: + // Inspect batch + log.Println(attempt.Batch.Entries()[0].Stmt) + } + + // Inspect or modify context + ctx = context.WithValue(ctx, "trace-id", "123") + + // Optionally bypass the handler and return an error to prevent query execution. + // For example, to simulate query timeouts. + if q.injectFault && attempt.Attempts == 0 { + <-time.After(1 * time.Second) + return nil, gocql.RequestErrWriteTimeout{} + } + + // The interceptor *must* invoke the handler to execute the query. + return handler(ctx) +} + +// Example_interceptor demonstrates how to implement a ExecAttemptInterceptor. +func Example_interceptor() { + cluster := gocql.NewCluster("localhost:9042") + cluster.ExecAttemptInterceptor = MyExecAttemptInterceptor{injectFault: true} + + 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.Fatalf("query failed %T", err) + } +} + +// Example_interceptor_chain demonstrates how to chain ExecAttemptInterceptors. +func Example_interceptor_chain() { + cluster := gocql.NewCluster("localhost:9042") + cluster.ExecAttemptInterceptor = gocql.ExecAttemptInterceptorChain{ + []gocql.ExecAttemptInterceptor{ + MyExecAttemptInterceptor{}, + MyExecAttemptInterceptor{}, + MyExecAttemptInterceptor{}, + }, + } + + 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.Fatalf("query failed %T", err) + } +} diff --git a/query_executor.go b/query_executor.go index 5d848c2e8..0d2ace1a7 100644 --- a/query_executor.go +++ b/query_executor.go @@ -26,6 +26,7 @@ package gocql import ( "context" + "net" "sync" "sync/atomic" "time" @@ -58,7 +59,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 @@ -69,16 +71,151 @@ type internalRequest interface { } type queryExecutor struct { - pool *policyConnPool - policy HostSelectionPolicy + pool *policyConnPool + policy HostSelectionPolicy + interceptor ExecAttemptInterceptor } -func (q *queryExecutor) attemptQuery(ctx context.Context, qry internalRequest, conn *Conn) *Iter { +type OpType int + +const ( + OpQuery OpType = iota + OpBatch +) + +type ImmutableQuery interface { + Statement() string + Values() []interface{} +} + +type immutableQuery struct { + query *Query +} + +func (q *immutableQuery) Statement() string { + return q.query.stmt +} + +func (q *immutableQuery) Values() []interface{} { + return q.query.values +} + +type ImmutableBatch interface { + Type() BatchType + Entries() []BatchEntry + Cons() Consistency +} + +type immutableBatch struct { + batch *Batch +} + +func (b *immutableBatch) Type() BatchType { + return b.batch.Type +} + +func (b *immutableBatch) Entries() []BatchEntry { + return b.batch.Entries +} + +func (b *immutableBatch) Cons() Consistency { + return b.batch.Cons +} + +type QueryAttempt struct { + // The op type, whether query or batch. + Type OpType + // Only one of Query or Batch will be set, depending on the op type. + Query ImmutableQuery + Batch ImmutableBatch + // The host that will receive the query. + Host *HostInfo + // The local address of the connection used to execute the query. + LocalAddr net.Addr + // The remote address of the connection used to execute the query. + RemoteAddr net.Addr + // The number of this query attempt. 0 for the initial attempt, 1 for the first retry, etc. Note that there may be multiple serial Attempts to different hosts within a single execution, whether main execution or speculative. + Attempts int + // The index of the speculative execution attempt this attempt is associated with, starting with index 1. -1 indicates this is the "main" execution. + SpeculativeExecutionCount int +} + +// QueryAttemptHandler is a function that attempts query execution. The interceptor must call this once if it does not return an error. +type QueryAttemptHandler = func(context.Context) (*Iter, error) + +// ExecAttemptInterceptor 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 ExecAttemptInterceptor interface { + // Intercept is invoked once immediately before a query or batch execution attempt, including retry attempts and + // speculative execution attempts. + // + // The interceptor is responsible for calling the `handler` function and returning the handler result. If the + // interceptor wants to bypass the handler and skip query execution, it should return an error. Failure to + // return either the handler result or an error will panic. + // + // Note that there is no affordance to mutate the query or batch at this stage -- the handler already encapsulates the original query/batch and cannot be modified. + Intercept(ctx context.Context, attempt QueryAttempt, handler QueryAttemptHandler) (*Iter, error) +} + +// Use an ExecAttemptInterceptorChain to apply multiple Interceptors at Intercept time. +type ExecAttemptInterceptorChain struct { + Interceptors []ExecAttemptInterceptor +} + +func (c ExecAttemptInterceptorChain) Intercept( + ctx context.Context, + attempt QueryAttempt, + handler QueryAttemptHandler, +) (*Iter, error) { + return c.Interceptors[0].Intercept(ctx, attempt, c.getNextHandler(0, attempt, handler)) +} + +func (c ExecAttemptInterceptorChain) getNextHandler(curr int, attempt QueryAttempt, final QueryAttemptHandler) QueryAttemptHandler { + if curr == len(c.Interceptors)-1 { + return final + } + + return func(ctx context.Context) (*Iter, error) { + return c.Interceptors[curr+1].Intercept(ctx, attempt, c.getNextHandler(curr+1, attempt, final)) + } +} + +func (q *queryExecutor) attemptQuery(ctx context.Context, iRequest internalRequest, conn *Conn, speculativeExecutionCount int) *Iter { start := time.Now() - iter := qry.execute(ctx, conn) - end := time.Now() - qry.attempt(q.pool.keyspace, end, start, iter, conn.host) + var iter *Iter + var err error + attempt := iRequest.getNextAttempt() + if q.interceptor != nil { + iq := QueryAttempt{ + Host: conn.host, + LocalAddr: conn.r.LocalAddr(), + RemoteAddr: conn.r.RemoteAddr(), + Attempts: attempt, + SpeculativeExecutionCount: speculativeExecutionCount, + } + if iQuery, ok := iRequest.(*internalQuery); ok { + iq.Type = OpQuery + iq.Query = &immutableQuery{query: iQuery.originalQuery} + } + if iBatch, ok := iRequest.(*internalBatch); ok { + iq.Type = OpBatch + iq.Batch = &immutableBatch{batch: iBatch.originalBatch} + } + iter, err = q.interceptor.Intercept(ctx, iq, func(cxCtx context.Context) (*Iter, error) { + it := iRequest.execute(cxCtx, conn) + return it, it.err + }) + if err != nil { + iter = &Iter{err: err} + } + } else { + iter = iRequest.execute(ctx, conn) + } + + end := time.Now() + iRequest.recordAttempt(attempt, q.pool.keyspace, end, start, iter, conn.host) return iter } @@ -91,7 +228,7 @@ func (q *queryExecutor) speculate(ctx context.Context, qry internalRequest, sp S for i := 0; i < sp.Attempts(); i++ { select { case <-ticker.C: - go q.run(ctx, qry, hostIter, results) + go q.run(ctx, qry, hostIter, i+1, results) case <-ctx.Done(): return newErrIter(ctx.Err(), qry.getQueryMetrics(), qry.Keyspace(), qry.getRoutingInfo(), qry.getKeyspaceFunc()) case iter := <-results: @@ -131,7 +268,7 @@ func (q *queryExecutor) executeQuery(qry internalRequest) (*Iter, error) { // it is, we force the policy to NonSpeculative sp := qry.speculativeExecutionPolicy() if qry.GetHostID() != "" || !qry.IsIdempotent() || sp.Attempts() == 0 { - return q.do(qry.Context(), qry, hostIter), nil + return q.do(qry.Context(), qry, hostIter, -1), nil } // When speculative execution is enabled, we could be accessing the host iterator from multiple goroutines below. @@ -150,7 +287,7 @@ func (q *queryExecutor) executeQuery(qry internalRequest) (*Iter, error) { results := make(chan *Iter, 1) // Launch the main execution - go q.run(ctx, qry, hostIter, results) + go q.run(ctx, qry, hostIter, -1, results) // The speculative executions are launched _in addition_ to the main // execution, on a timer. So Speculation{2} would make 3 executions running @@ -167,7 +304,7 @@ func (q *queryExecutor) executeQuery(qry internalRequest) (*Iter, error) { } } -func (q *queryExecutor) do(ctx context.Context, qry internalRequest, hostIter NextHost) *Iter { +func (q *queryExecutor) do(ctx context.Context, qry internalRequest, hostIter NextHost, speculativeExecutionCount int) *Iter { selectedHost := hostIter() rt := qry.retryPolicy() @@ -192,7 +329,7 @@ func (q *queryExecutor) do(ctx context.Context, qry internalRequest, hostIter Ne continue } - iter = q.attemptQuery(ctx, qry, conn) + iter = q.attemptQuery(ctx, qry, conn, speculativeExecutionCount) iter.host = selectedHost.Info() // Update host switch iter.err { @@ -248,9 +385,9 @@ func (q *queryExecutor) do(ctx context.Context, qry internalRequest, hostIter Ne return newErrIter(ErrNoConnections, qry.getQueryMetrics(), qry.Keyspace(), qry.getRoutingInfo(), qry.getKeyspaceFunc()) } -func (q *queryExecutor) run(ctx context.Context, qry internalRequest, hostIter NextHost, results chan<- *Iter) { +func (q *queryExecutor) run(ctx context.Context, qry internalRequest, hostIter NextHost, speculativeExecutionCount int, results chan<- *Iter) { select { - case results <- q.do(ctx, qry, hostIter): + case results <- q.do(ctx, qry, hostIter, speculativeExecutionCount): case <-ctx.Done(): } } @@ -371,14 +508,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(attemptNum, latency, q.session) if q.qryOpts.observer != nil { metricsForHost := q.hostMetricsManager.attempt(latency, host) @@ -397,7 +538,7 @@ func (q *internalQuery) attempt(keyspace string, end, start time.Time, iter *Ite Host: host, Metrics: metricsForHost, Err: iter.err, - Attempt: attempt, + Attempt: attemptNum, Query: q.originalQuery, }) } @@ -601,14 +742,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(attemptNum, latency, b.session) if b.batchOpts.observer == nil { return @@ -644,7 +789,7 @@ func (b *internalBatch) attempt(keyspace string, end, start time.Time, iter *Ite Host: host, Metrics: metricsForHost, Err: iter.err, - Attempt: attempt, + Attempt: attemptNum, Batch: b.originalBatch, }) } diff --git a/session.go b/session.go index dee2a8fe1..8dc00e02a 100644 --- a/session.go +++ b/session.go @@ -230,8 +230,9 @@ func NewSession(cfg ClusterConfig) (*Session, error) { // set the executor here in case the policy needs to execute queries in Init s.executor = &queryExecutor{ - pool: s.pool, - policy: cfg.PoolConfig.HostSelectionPolicy, + pool: s.pool, + policy: cfg.PoolConfig.HostSelectionPolicy, + interceptor: cfg.ExecAttemptInterceptor, } s.policy.Init(s) @@ -968,26 +969,49 @@ type hostMetrics struct { TotalLatency int64 } +// Query attempts could be started and finished out of order when using speculative execution. Therefore, +// we issue attempt indexes at query start time (so that Interceptors can use them) then record those attempts +// at completion time (so that metrics can correctly express average latency). 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(attempt int, addLatency time.Duration, s *Session) { + qm.l.Lock() + defer qm.l.Unlock() + qm.totalLatency += addLatency.Nanoseconds() + qm.attemptsCompleted++ + if attempt > qm.attemptsCompleted && s != nil { + s.logger.Debug("attempt number is greater than total attempts completed, this should not happen", NewLogFieldInt("attempt", attempt), NewLogFieldInt("attemptsCompleted", qm.attemptsCompleted)) + } } +// Returns total number of attempts started. func (qm *queryMetrics) attempts() int { - return int(atomic.LoadInt64(&qm.totalAttempts)) + qm.l.RLock() + defer qm.l.RUnlock() + return qm.attemptsStarted } 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 {