fingerprint_builtins.go

// Copyright 2023 The Cockroach Authors.
//
// Use of this software is governed by the Business Source License
// included in the file licenses/BSL.txt.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0, included in the file
// licenses/APL.txt.

package builtins

import (
	"context"
	"fmt"
	"time"

	"github.com/cockroachdb/cockroach/pkg/clusterversion"
	"github.com/cockroachdb/cockroach/pkg/kv"
	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/lock"
	"github.com/cockroachdb/cockroach/pkg/roachpb"
	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
	"github.com/cockroachdb/cockroach/pkg/storage"
	"github.com/cockroachdb/cockroach/pkg/util/admission/admissionpb"
	"github.com/cockroachdb/cockroach/pkg/util/contextutil"
	"github.com/cockroachdb/cockroach/pkg/util/hlc"
	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
	"github.com/cockroachdb/cockroach/pkg/util/tracing"
	"github.com/cockroachdb/errors"
)

func fingerprint(
	ctx context.Context,
	evalCtx *eval.Context,
	args tree.Datums,
	startTime hlc.Timestamp,
	allRevisions, stripped bool,
) (tree.Datum, error) {
	ctx, sp := tracing.ChildSpan(ctx, "crdb_internal.fingerprint")
	defer sp.Finish()

	if !evalCtx.Settings.Version.IsActive(ctx, clusterversion.V23_1) {
		return nil, errors.Errorf("cannot use crdb_internal.fingerprint until the cluster version is at least %s",
			clusterversion.V23_1.String())
	}

	isAdmin, err := evalCtx.SessionAccessor.HasAdminRole(ctx)
	if err != nil {
		return nil, err
	}
	if !isAdmin {
		return nil, errors.New("crdb_internal.fingerprint() requires admin privilege")
	}

	arr := tree.MustBeDArray(args[0])
	if arr.Len() != 2 {
		return nil, errors.New("expected an array of two elements")
	}
	startKey := []byte(tree.MustBeDBytes(arr.Array[0]))
	endKey := []byte(tree.MustBeDBytes(arr.Array[1]))

	filter := kvpb.MVCCFilter_Latest
	if allRevisions {
		filter = kvpb.MVCCFilter_All
	}
	req := &kvpb.ExportRequest{
		RequestHeader:      kvpb.RequestHeader{Key: startKey, EndKey: endKey},
		StartTime:          startTime,
		MVCCFilter:         filter,
		ExportFingerprint:  true,
		FingerprintOptions: kvpb.FingerprintOptions{Stripped: stripped}}
	header := kvpb.Header{
		Timestamp: evalCtx.Txn.ReadTimestamp(),
		// We set WaitPolicy to Error, so that the export will return an error
		// to us instead of a blocking wait if it hits any other txns.
		//
		// TODO(adityamaru): We might need to handle WriteIntentErrors
		// specially in the future so as to allow the fingerprint to complete
		// in the face of intents.
		WaitPolicy: lock.WaitPolicy_Error,
		// TODO(ssd): Setting this disables async sending in
		// DistSender so it likely substantially impacts
		// performance.
		ReturnElasticCPUResumeSpans: true,
	}
	admissionHeader := kvpb.AdmissionHeader{
		Priority:                 int32(admissionpb.BulkNormalPri),
		CreateTime:               timeutil.Now().UnixNano(),
		Source:                   kvpb.AdmissionHeader_FROM_SQL,
		NoMemoryReservedAtSource: true,
	}

	todo := make(chan *kvpb.ExportRequest, 1)
	todo <- req
	ctxDone := ctx.Done()
	var fingerprint uint64
	// TODO(adityamaru): Memory monitor this slice of buffered SSTs that
	// contain range keys across ExportRequests.
	ssts := make([][]byte, 0)
	for {
		select {
		case <-ctxDone:
			return nil, ctx.Err()
		case req := <-todo:
			var rawResp kvpb.Response
			var pErr *kvpb.Error
			exportRequestErr := contextutil.RunWithTimeout(ctx,
				fmt.Sprintf("ExportRequest fingerprint for span %s", roachpb.Span{Key: startKey, EndKey: endKey}),
				5*time.Minute, func(ctx context.Context) error {
					rawResp, pErr = kv.SendWrappedWithAdmission(ctx,
						evalCtx.Txn.DB().NonTransactionalSender(), header, admissionHeader, req)
					if pErr != nil {
						return pErr.GoError()
					}
					return nil
				})
			if exportRequestErr != nil {
				return nil, exportRequestErr
			}

			resp := rawResp.(*kvpb.ExportResponse)
			for _, file := range resp.Files {
				fingerprint = fingerprint ^ file.Fingerprint

				// Aggregate all the range keys that need fingerprinting once all
				// ExportRequests have been completed.
				if len(file.SST) != 0 {
					ssts = append(ssts, file.SST)
				}
			}
			if resp.ResumeSpan != nil {
				if !resp.ResumeSpan.Valid() {
					return nil, errors.Errorf("invalid resume span: %s", resp.ResumeSpan)
				}

				resumeReq := req
				resumeReq.RequestHeader = kvpb.RequestHeaderFromSpan(*resp.ResumeSpan)
				todo <- resumeReq
			}
		default:
			// No ExportRequests left to send. We've aggregated range keys
			// across all ExportRequests and can now fingerprint them.
			//
			// NB: We aggregate rangekeys across ExportRequests and then
			// fingerprint them on the client, instead of fingerprinting them as
			// part of the ExportRequest command evaluation, because range keys
			// do not have a stable, discrete identity. Their fragmentation can
			// be influenced by rangekeys outside the time interval that we are
			// fingerprinting, or by range splits. So, we need to "defragment"
			// all the rangekey stacks we observe such that the fragmentation is
			// deterministic on only the data we want to fingerprint in our key
			// and time interval.
			//
			// Egs:
			//
			// t2  				[-----)[----)
			//
			// t1 	[----)[-----)
			//			a			b			c			d
			//
			// Assume we have two rangekeys [a, c)@t1 and [b, d)@t2. They will
			// fragment as shown in the diagram above. If we wish to fingerprint
			// key [a-d) in time interval (t1, t2] the fragmented rangekey
			// [a, c)@t1 is outside our time interval and should not influence our
			// fingerprint. The iterator in `fingerprintRangekeys` will
			// "defragment" the rangekey stacks [b-c)@t2 and [c-d)@t2 and
			// fingerprint them as a single rangekey with bounds [b-d)@t2.
			rangekeyFingerprint, err := storage.FingerprintRangekeys(ctx, evalCtx.Settings,
				storage.MVCCExportFingerprintOptions{
					StripTenantPrefix:  true,
					StripValueChecksum: true,
					StrippedVersion:    stripped,
				}, ssts)
			if err != nil {
				return nil, err
			}
			fingerprint = fingerprint ^ rangekeyFingerprint
			return tree.NewDInt(tree.DInt(fingerprint)), nil
		}
	}
}