github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/kv/kvserver/batcheval/command.go (about)

     1  // Copyright 2014 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 batcheval
    12  
    13  import (
    14  	"context"
    15  
    16  	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/batcheval/result"
    17  	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/spanset"
    18  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    19  	"github.com/cockroachdb/cockroach/pkg/storage"
    20  	"github.com/cockroachdb/cockroach/pkg/util/log"
    21  )
    22  
    23  // declareKeysFunc adds all key spans that a command touches to the latchSpans
    24  // set. It then adds all key spans within which that the command expects to have
    25  // isolation from conflicting transactions to the lockSpans set.
    26  type declareKeysFunc func(
    27  	_ *roachpb.RangeDescriptor, _ roachpb.Header, _ roachpb.Request, latchSpans, lockSpans *spanset.SpanSet,
    28  )
    29  
    30  // A Command is the implementation of a single request within a BatchRequest.
    31  type Command struct {
    32  	// DeclareKeys adds all keys this command touches, and when (if applicable),
    33  	// to the given SpanSet.
    34  	//
    35  	// TODO(nvanbenschoten): rationalize this RangeDescriptor. Can it change
    36  	// between key declaration and cmd evaluation? Really, do it.
    37  	DeclareKeys declareKeysFunc
    38  
    39  	// Eval{RW,RO} evaluates a read-{write,only} command respectively on the
    40  	// given engine.{ReadWriter,Reader}. This is typically derived from
    41  	// engine.NewBatch or engine.NewReadOnly (which is more performant than
    42  	// engine.Batch for read-only commands).
    43  	// It should populate the supplied response (always a non-nil pointer to the
    44  	// correct type) and return special side effects (if any) in the Result. If
    45  	// it writes to the engine it should also update *CommandArgs.Stats. It
    46  	// should treat the provided request as immutable.
    47  	//
    48  	// Only one of these is ever set at a time.
    49  	EvalRW func(context.Context, storage.ReadWriter, CommandArgs, roachpb.Response) (result.Result, error)
    50  	EvalRO func(context.Context, storage.Reader, CommandArgs, roachpb.Response) (result.Result, error)
    51  }
    52  
    53  var cmds = make(map[roachpb.Method]Command)
    54  
    55  // RegisterReadWriteCommand makes a read-write command available for execution.
    56  // It must only be called before any evaluation takes place.
    57  func RegisterReadWriteCommand(
    58  	method roachpb.Method,
    59  	declare declareKeysFunc,
    60  	impl func(context.Context, storage.ReadWriter, CommandArgs, roachpb.Response) (result.Result, error),
    61  ) {
    62  	register(method, Command{
    63  		DeclareKeys: declare,
    64  		EvalRW:      impl,
    65  	})
    66  }
    67  
    68  // RegisterReadOnlyCommand makes a read-only command available for execution. It
    69  // must only be called before any evaluation takes place.
    70  func RegisterReadOnlyCommand(
    71  	method roachpb.Method,
    72  	declare declareKeysFunc,
    73  	impl func(context.Context, storage.Reader, CommandArgs, roachpb.Response) (result.Result, error),
    74  ) {
    75  	register(method, Command{
    76  		DeclareKeys: declare,
    77  		EvalRO:      impl,
    78  	})
    79  }
    80  
    81  func register(method roachpb.Method, command Command) {
    82  	if _, ok := cmds[method]; ok {
    83  		log.Fatalf(context.TODO(), "cannot overwrite previously registered method %v", method)
    84  	}
    85  	cmds[method] = command
    86  }
    87  
    88  // UnregisterCommand is provided for testing and allows removing a command.
    89  // It is a no-op if the command is not registered.
    90  func UnregisterCommand(method roachpb.Method) {
    91  	delete(cmds, method)
    92  }
    93  
    94  // LookupCommand returns the command for the given method, with the boolean
    95  // indicating success or failure.
    96  func LookupCommand(method roachpb.Method) (Command, bool) {
    97  	cmd, ok := cmds[method]
    98  	return cmd, ok
    99  }