github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/sqlbase/cancel_checker.go (about)

     1  // Copyright 2017 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package sqlbase
    12  
    13  import (
    14  	"context"
    15  	"sync/atomic"
    16  )
    17  
    18  // CancelChecker is a helper object for repeatedly checking whether the associated context
    19  // has been canceled or not. Encapsulates all logic for waiting for cancelCheckInterval
    20  // rows before actually checking for cancellation. The cancellation check
    21  // has a significant time overhead, so it's not checked in every iteration.
    22  type CancelChecker struct {
    23  	// Reference to associated context to check.
    24  	ctx context.Context
    25  
    26  	// Number of times Check() has been called since last context cancellation check.
    27  	callsSinceLastCheck uint32
    28  }
    29  
    30  // NewCancelChecker returns a new CancelChecker.
    31  func NewCancelChecker(ctx context.Context) *CancelChecker {
    32  	return &CancelChecker{
    33  		ctx: ctx,
    34  	}
    35  }
    36  
    37  // Check returns an error if the associated query has been canceled.
    38  func (c *CancelChecker) Check() error {
    39  	// Interval of Check() calls to wait between checks for context
    40  	// cancellation. The value is a power of 2 to allow the compiler to use
    41  	// bitwise AND instead of division.
    42  	const cancelCheckInterval = 1024
    43  
    44  	if atomic.LoadUint32(&c.callsSinceLastCheck)%cancelCheckInterval == 0 {
    45  		select {
    46  		case <-c.ctx.Done():
    47  			// Once the context is canceled, we no longer increment
    48  			// callsSinceLastCheck and will fall into this path on subsequent calls
    49  			// to Check().
    50  			return QueryCanceledError
    51  		default:
    52  		}
    53  	}
    54  
    55  	// Increment. This may rollover when the 32-bit capacity is reached,
    56  	// but that's all right.
    57  	atomic.AddUint32(&c.callsSinceLastCheck, 1)
    58  	return nil
    59  }
    60  
    61  // Reset resets this cancel checker with a fresh context.
    62  func (c *CancelChecker) Reset(ctx context.Context) {
    63  	*c = CancelChecker{
    64  		ctx: ctx,
    65  	}
    66  }