github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/testutils/kvclientutils/txn_recovery.go (about)

     1  // Copyright 2020 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 kvclientutils
    12  
    13  import (
    14  	"context"
    15  	"fmt"
    16  
    17  	"github.com/cockroachdb/cockroach/pkg/kv"
    18  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    19  	"github.com/cockroachdb/cockroach/pkg/util/hlc"
    20  	"github.com/cockroachdb/cockroach/pkg/util/tracing"
    21  	"github.com/cockroachdb/errors"
    22  )
    23  
    24  // PushExpectation expresses an expectation for CheckPushResult about what the
    25  // push did.
    26  type PushExpectation int
    27  
    28  const (
    29  	// ExpectPusheeTxnRecovery means we're expecting transaction recovery to be
    30  	// performed (after finding a STAGING txn record).
    31  	ExpectPusheeTxnRecovery PushExpectation = iota
    32  	// ExpectPusheeTxnRecordNotFound means we're expecting the push to not find the
    33  	// pushee txn record.
    34  	ExpectPusheeTxnRecordNotFound
    35  	// DontExpectAnything means we're not going to check the state in which the
    36  	// pusher found the pushee's txn record.
    37  	DontExpectAnything
    38  )
    39  
    40  // ExpectedTxnResolution expresses an expectation for CheckPushResult about the
    41  // outcome of the push.
    42  type ExpectedTxnResolution int
    43  
    44  const (
    45  	// ExpectAborted means that the pushee is expected to have been aborted. Note
    46  	// that a committed txn that has been cleaned up also results in an ABORTED
    47  	// result for a pusher.
    48  	ExpectAborted ExpectedTxnResolution = iota
    49  	// ExpectCommitted means that the pushee is expected to have found the pushee
    50  	// to be committed - or STAGING in which case the push will have performed
    51  	// successful transaction recovery.
    52  	ExpectCommitted
    53  )
    54  
    55  // CheckPushResult pushes the specified txn and checks that the pushee's
    56  // resolution is the expected one.
    57  func CheckPushResult(
    58  	ctx context.Context,
    59  	db *kv.DB,
    60  	txn roachpb.Transaction,
    61  	expResolution ExpectedTxnResolution,
    62  	pushExpectation PushExpectation,
    63  ) error {
    64  	pushReq := roachpb.PushTxnRequest{
    65  		RequestHeader: roachpb.RequestHeader{
    66  			Key: txn.Key,
    67  		},
    68  		PusheeTxn: txn.TxnMeta,
    69  		PushTo:    hlc.Timestamp{},
    70  		PushType:  roachpb.PUSH_ABORT,
    71  		// We're going to Force the push in order to not wait for the pushee to
    72  		// expire.
    73  		Force: true,
    74  	}
    75  	ba := roachpb.BatchRequest{}
    76  	ba.Add(&pushReq)
    77  
    78  	recCtx, collectRec, cancel := tracing.ContextWithRecordingSpan(ctx, "test trace")
    79  	defer cancel()
    80  
    81  	resp, pErr := db.NonTransactionalSender().Send(recCtx, ba)
    82  	if pErr != nil {
    83  		return pErr.GoError()
    84  	}
    85  
    86  	var statusErr error
    87  	pusheeStatus := resp.Responses[0].GetPushTxn().PusheeTxn.Status
    88  	switch pusheeStatus {
    89  	case roachpb.ABORTED:
    90  		if expResolution != ExpectAborted {
    91  			statusErr = errors.Errorf("transaction unexpectedly aborted")
    92  		}
    93  	case roachpb.COMMITTED:
    94  		if expResolution != ExpectCommitted {
    95  			statusErr = errors.Errorf("transaction unexpectedly committed")
    96  		}
    97  	default:
    98  		return errors.Errorf("unexpected txn status: %s", pusheeStatus)
    99  	}
   100  
   101  	// Verify that we're not fooling ourselves and that checking for the implicit
   102  	// commit actually caused the txn recovery procedure to run.
   103  	recording := collectRec()
   104  	var resolutionErr error
   105  	switch pushExpectation {
   106  	case ExpectPusheeTxnRecovery:
   107  		expMsg := fmt.Sprintf("recovered txn %s", txn.ID.Short())
   108  		if _, ok := recording.FindLogMessage(expMsg); !ok {
   109  			resolutionErr = errors.Errorf(
   110  				"recovery didn't run as expected (missing \"%s\"). recording: %s",
   111  				expMsg, recording)
   112  		}
   113  	case ExpectPusheeTxnRecordNotFound:
   114  		expMsg := "pushee txn record not found"
   115  		if _, ok := recording.FindLogMessage(expMsg); !ok {
   116  			resolutionErr = errors.Errorf(
   117  				"push didn't run as expected (missing \"%s\"). recording: %s",
   118  				expMsg, recording)
   119  		}
   120  	case DontExpectAnything:
   121  	}
   122  
   123  	return errors.CombineErrors(statusErr, resolutionErr)
   124  }