diff --git a/executor.go b/executor.go index dca0bd546..bbb8c2297 100644 --- a/executor.go +++ b/executor.go @@ -87,6 +87,11 @@ func (e *executor) Execute(ctx context.Context, index string, q *pql.Query, shar resp := QueryResponse{} + // Check for query cancellation. + if err := validateQueryContext(ctx); err != nil { + return resp, err + } + // Verify that an index is set. if index == "" { return resp, ErrIndexRequired @@ -112,12 +117,16 @@ func (e *executor) Execute(ctx context.Context, index string, q *pql.Query, shar if !opt.Remote { if err := e.translateCalls(ctx, index, idx, q.Calls); err != nil { return resp, err + } else if err := validateQueryContext(ctx); err != nil { + return resp, err } } results, err := e.execute(ctx, index, q, shards, opt) if err != nil { return resp, err + } else if err := validateQueryContext(ctx); err != nil { + return resp, err } resp.Results = results @@ -159,6 +168,8 @@ func (e *executor) Execute(ctx context.Context, index string, q *pql.Query, shar if !opt.Remote { if err := e.translateResults(ctx, index, idx, q.Calls, results); err != nil { return resp, err + } else if err := validateQueryContext(ctx); err != nil { + return resp, err } } @@ -217,6 +228,10 @@ func (e *executor) execute(ctx context.Context, index string, q *pql.Query, shar // Execute each call serially. results := make([]interface{}, 0, len(q.Calls)) for _, call := range q.Calls { + if err := validateQueryContext(ctx); err != nil { + return nil, err + } + v, err := e.executeCall(ctx, index, call, shards, opt) if err != nil { return nil, err @@ -231,7 +246,9 @@ func (e *executor) executeCall(ctx context.Context, index string, c *pql.Call, s span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeCall") defer span.Finish() - if err := e.validateCallArgs(c); err != nil { + if err := validateQueryContext(ctx); err != nil { + return nil, err + } else if err := e.validateCallArgs(c); err != nil { return nil, errors.Wrap(err, "validating args") } indexTag := fmt.Sprintf("index:%s", index) @@ -521,6 +538,10 @@ func (e *executor) executeBitmapCall(ctx context.Context, index string, c *pql.C // executeBitmapCallShard executes a bitmap call for a single shard. func (e *executor) executeBitmapCallShard(ctx context.Context, index string, c *pql.Call, shard uint64) (*Row, error) { + if err := validateQueryContext(ctx); err != nil { + return nil, err + } + span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeBitmapCallShard") defer span.Finish() @@ -1977,7 +1998,13 @@ func (e *executor) executeBulkSetRowAttrs(ctx context.Context, index string, cal // Collect attributes by field/id. m := make(map[string]map[uint64]map[string]interface{}) - for _, c := range calls { + for i, c := range calls { + if i%10 == 0 { + if err := validateQueryContext(ctx); err != nil { + return nil, err + } + } + field, ok := c.Args["_field"].(string) if !ok { return nil, errors.New("SetRowAttrs() field required") @@ -2560,6 +2587,23 @@ func (e *executor) translateResult(index string, idx *Index, call *pql.Call, res return result, nil } +// validateQueryContext returns a query-appropriate error if the context is done. +func validateQueryContext(ctx context.Context) error { + select { + case <-ctx.Done(): + switch err := ctx.Err(); err { + case context.Canceled: + return ErrQueryCancelled + case context.DeadlineExceeded: + return ErrQueryTimeout + default: + return err + } + default: + return nil + } +} + // errShardUnavailable is a marker error if no nodes are available. var errShardUnavailable = errors.New("shard unavailable") diff --git a/pilosa.go b/pilosa.go index d2fdf8144..56e8e0f9e 100644 --- a/pilosa.go +++ b/pilosa.go @@ -54,6 +54,8 @@ var ( // ErrFragmentNotFound is returned when a fragment does not exist. ErrFragmentNotFound = errors.New("fragment not found") ErrQueryRequired = errors.New("query required") + ErrQueryCancelled = errors.New("query cancelled") + ErrQueryTimeout = errors.New("query timeout") ErrTooManyWrites = errors.New("too many write commands") ErrClusterDoesNotOwnShard = errors.New("cluster does not own shard")