github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/ccl/storageccl/writebatch.go (about)

     1  // Copyright 2016 The Cockroach Authors.
     2  //
     3  // Licensed as a CockroachDB Enterprise file under the Cockroach Community
     4  // License (the "License"); you may not use this file except in compliance with
     5  // the License. You may obtain a copy of the License at
     6  //
     7  //     https://github.com/cockroachdb/cockroach/blob/master/licenses/CCL.txt
     8  
     9  package storageccl
    10  
    11  import (
    12  	"context"
    13  	"fmt"
    14  
    15  	"github.com/cockroachdb/cockroach/pkg/ccl/storageccl/engineccl"
    16  	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/batcheval"
    17  	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/batcheval/result"
    18  	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/spanset"
    19  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    20  	"github.com/cockroachdb/cockroach/pkg/storage"
    21  	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
    22  	"github.com/cockroachdb/cockroach/pkg/util/log"
    23  	"github.com/cockroachdb/cockroach/pkg/util/tracing"
    24  	"github.com/cockroachdb/errors"
    25  )
    26  
    27  func init() {
    28  	batcheval.RegisterReadWriteCommand(roachpb.WriteBatch, batcheval.DefaultDeclareKeys, evalWriteBatch)
    29  }
    30  
    31  // evalWriteBatch applies the operations encoded in a BatchRepr. Any existing
    32  // data in the affected keyrange is first cleared (not tombstoned), which makes
    33  // this command idempotent.
    34  func evalWriteBatch(
    35  	ctx context.Context, batch storage.ReadWriter, cArgs batcheval.CommandArgs, _ roachpb.Response,
    36  ) (result.Result, error) {
    37  
    38  	args := cArgs.Args.(*roachpb.WriteBatchRequest)
    39  	h := cArgs.Header
    40  	ms := cArgs.Stats
    41  
    42  	_, span := tracing.ChildSpan(ctx, fmt.Sprintf("WriteBatch [%s,%s)", args.Key, args.EndKey))
    43  	defer tracing.FinishSpan(span)
    44  	if log.V(1) {
    45  		log.Infof(ctx, "writebatch [%s,%s)", args.Key, args.EndKey)
    46  	}
    47  
    48  	// We can't use the normal RangeKeyMismatchError mechanism for dealing with
    49  	// splits because args.Data should stay an opaque blob to DistSender.
    50  	if args.DataSpan.Key.Compare(args.Key) < 0 || args.DataSpan.EndKey.Compare(args.EndKey) > 0 {
    51  		// TODO(dan): Add a new field in roachpb.Error, so the client can catch
    52  		// this and retry.
    53  		return result.Result{}, errors.New("data spans multiple ranges")
    54  	}
    55  
    56  	mvccStartKey := storage.MVCCKey{Key: args.Key}
    57  	mvccEndKey := storage.MVCCKey{Key: args.EndKey}
    58  
    59  	// Verify that the keys in the batch are within the range specified by the
    60  	// request header.
    61  	msBatch, err := engineccl.VerifyBatchRepr(args.Data, mvccStartKey, mvccEndKey, h.Timestamp.WallTime)
    62  	if err != nil {
    63  		return result.Result{}, err
    64  	}
    65  	ms.Add(msBatch)
    66  
    67  	// Check if there was data in the affected keyrange. If so, delete it (and
    68  	// adjust the MVCCStats) before applying the WriteBatch data.
    69  	existingStats, err := clearExistingData(ctx, batch, args.Key, args.EndKey, h.Timestamp.WallTime)
    70  	if err != nil {
    71  		return result.Result{}, errors.Wrap(err, "clearing existing data")
    72  	}
    73  	ms.Subtract(existingStats)
    74  
    75  	if err := batch.ApplyBatchRepr(args.Data, false /* sync */); err != nil {
    76  		return result.Result{}, err
    77  	}
    78  	return result.Result{}, nil
    79  }
    80  
    81  func clearExistingData(
    82  	ctx context.Context, batch storage.ReadWriter, start, end roachpb.Key, nowNanos int64,
    83  ) (enginepb.MVCCStats, error) {
    84  	{
    85  		isEmpty := true
    86  		if err := batch.Iterate(start, end, func(_ storage.MVCCKeyValue) (bool, error) {
    87  			isEmpty = false
    88  			return true, nil // stop right away
    89  		}); err != nil {
    90  			return enginepb.MVCCStats{}, errors.Wrap(err, "while checking for empty key space")
    91  		}
    92  
    93  		if isEmpty {
    94  			return enginepb.MVCCStats{}, nil
    95  		}
    96  	}
    97  
    98  	iter := batch.NewIterator(storage.IterOptions{UpperBound: end})
    99  	defer iter.Close()
   100  
   101  	iter.SeekGE(storage.MakeMVCCMetadataKey(start))
   102  	if ok, err := iter.Valid(); err != nil {
   103  		return enginepb.MVCCStats{}, err
   104  	} else if ok && !iter.UnsafeKey().Less(storage.MakeMVCCMetadataKey(end)) {
   105  		return enginepb.MVCCStats{}, nil
   106  	}
   107  
   108  	existingStats, err := iter.ComputeStats(start, end, nowNanos)
   109  	if err != nil {
   110  		return enginepb.MVCCStats{}, err
   111  	}
   112  
   113  	log.Eventf(ctx, "target key range not empty, will clear existing data: %+v", existingStats)
   114  	// If this is a Iterator, we have to unwrap it because
   115  	// ClearIterRange needs a plain rocksdb iterator (and can't unwrap
   116  	// it itself because of import cycles).
   117  	if ssi, ok := iter.(*spanset.Iterator); ok {
   118  		iter = ssi.Iterator()
   119  	}
   120  	// TODO(dan): Ideally, this would use `batch.ClearRange` but it doesn't
   121  	// yet work with read-write batches (or IngestExternalData).
   122  	if err := batch.ClearIterRange(iter, start, end); err != nil {
   123  		return enginepb.MVCCStats{}, err
   124  	}
   125  	return existingStats, nil
   126  }