github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/kv/kvserver/batcheval/cmd_heartbeat_txn.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  	"fmt"
    16  
    17  	"github.com/cockroachdb/cockroach/pkg/keys"
    18  	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/batcheval/result"
    19  	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/spanset"
    20  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    21  	"github.com/cockroachdb/cockroach/pkg/storage"
    22  	"github.com/cockroachdb/cockroach/pkg/util/hlc"
    23  )
    24  
    25  func init() {
    26  	RegisterReadWriteCommand(roachpb.HeartbeatTxn, declareKeysHeartbeatTransaction, HeartbeatTxn)
    27  }
    28  
    29  func declareKeysHeartbeatTransaction(
    30  	desc *roachpb.RangeDescriptor,
    31  	header roachpb.Header,
    32  	req roachpb.Request,
    33  	latchSpans, _ *spanset.SpanSet,
    34  ) {
    35  	declareKeysWriteTransaction(desc, header, req, latchSpans)
    36  }
    37  
    38  // HeartbeatTxn updates the transaction status and heartbeat
    39  // timestamp after receiving transaction heartbeat messages from
    40  // coordinator. Returns the updated transaction.
    41  func HeartbeatTxn(
    42  	ctx context.Context, readWriter storage.ReadWriter, cArgs CommandArgs, resp roachpb.Response,
    43  ) (result.Result, error) {
    44  	args := cArgs.Args.(*roachpb.HeartbeatTxnRequest)
    45  	h := cArgs.Header
    46  	reply := resp.(*roachpb.HeartbeatTxnResponse)
    47  
    48  	if err := VerifyTransaction(h, args, roachpb.PENDING, roachpb.STAGING); err != nil {
    49  		return result.Result{}, err
    50  	}
    51  
    52  	if args.Now.IsEmpty() {
    53  		return result.Result{}, fmt.Errorf("now not specified for heartbeat")
    54  	}
    55  
    56  	key := keys.TransactionKey(h.Txn.Key, h.Txn.ID)
    57  
    58  	var txn roachpb.Transaction
    59  	if ok, err := storage.MVCCGetProto(
    60  		ctx, readWriter, key, hlc.Timestamp{}, &txn, storage.MVCCGetOptions{},
    61  	); err != nil {
    62  		return result.Result{}, err
    63  	} else if !ok {
    64  		// No existing transaction record was found - create one by writing
    65  		// it below.
    66  		txn = *h.Txn
    67  
    68  		// Verify that it is safe to create the transaction record.
    69  		if err := CanCreateTxnRecord(cArgs.EvalCtx, &txn); err != nil {
    70  			return result.Result{}, err
    71  		}
    72  	}
    73  
    74  	if !txn.Status.IsFinalized() {
    75  		// NOTE: this only updates the LastHeartbeat. It doesn't update any other
    76  		// field from h.Txn, even if it could. Whether that's a good thing or not
    77  		// is up for debate.
    78  		txn.LastHeartbeat.Forward(args.Now)
    79  		txnRecord := txn.AsRecord()
    80  		if err := storage.MVCCPutProto(ctx, readWriter, cArgs.Stats, key, hlc.Timestamp{}, nil, &txnRecord); err != nil {
    81  			return result.Result{}, err
    82  		}
    83  	}
    84  
    85  	reply.Txn = &txn
    86  	return result.Result{}, nil
    87  }