From 5c1e92b5a2dd21eff3a87d5a8ff5b27db67ed71e Mon Sep 17 00:00:00 2001
From: Oleg Afanasyev <oleg@cockroachlabs.com>
Date: Wed, 1 Sep 2021 09:33:27 +0100
Subject: [PATCH] kvserver: add cluster configuration to throttle exports

Export requests could iterate unbounded amount of data in storage
this diff adds kv.bulk_sst.max_request_time hidden cluster setting
to limit how for long export could run irrespective of how much data
is actually exported.

Release note: None
---
 pkg/kv/kvserver/batcheval/BUILD.bazel   |  1 +
 pkg/kv/kvserver/batcheval/cmd_export.go | 18 ++++++++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/pkg/kv/kvserver/batcheval/BUILD.bazel b/pkg/kv/kvserver/batcheval/BUILD.bazel
index 94b435889d1e..1d3a64f637b6 100644
--- a/pkg/kv/kvserver/batcheval/BUILD.bazel
+++ b/pkg/kv/kvserver/batcheval/BUILD.bazel
@@ -80,6 +80,7 @@ go_library(
         "//pkg/util/log",
         "//pkg/util/mon",
         "//pkg/util/protoutil",
+        "//pkg/util/timeutil",
         "//pkg/util/tracing",
         "//pkg/util/uuid",
         "@com_github_cockroachdb_errors//:errors",
diff --git a/pkg/kv/kvserver/batcheval/cmd_export.go b/pkg/kv/kvserver/batcheval/cmd_export.go
index 9cdc78652170..2ec7f7fb5a6a 100644
--- a/pkg/kv/kvserver/batcheval/cmd_export.go
+++ b/pkg/kv/kvserver/batcheval/cmd_export.go
@@ -21,6 +21,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/storage"
 	"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"
 	"github.com/gogo/protobuf/types"
@@ -58,6 +59,20 @@ var ExportRequestMaxAllowedFileSizeOverage = settings.RegisterByteSizeSetting(
 	64<<20, /* 64 MiB */
 ).WithPublic()
 
+// exportRequestMaxIterationTime controls time spent by export request iterating
+// over data in underlying storage. This threshold preventing export request from
+// holding locks for too long and preventing non mvcc operations from progressing.
+// If request takes longer than this threshold it would stop and return already
+// collected data and allow caller to use resume span to continue.
+var exportRequestMaxIterationTime = settings.RegisterDurationSetting(
+	"kv.bulk_sst.max_request_time",
+	"if set, limits amount of time spent in export requests; "+
+		"if export request can not finish within allocated time it will resume from the point it stopped in "+
+		"subsequent request",
+	// Feature is disabled by default.
+	0,
+)
+
 func init() {
 	RegisterReadOnlyCommand(roachpb.Export, declareKeysExport, evalExport)
 }
@@ -136,6 +151,8 @@ func evalExport(
 		maxSize = targetSize + uint64(allowedOverage)
 	}
 
+	maxRunTime := exportRequestMaxIterationTime.Get(&cArgs.EvalCtx.ClusterSettings().SV)
+
 	// Time-bound iterators only make sense to use if the start time is set.
 	useTBI := args.EnableTimeBoundIteratorOptimization && !args.StartTime.IsEmpty()
 	// Only use resume timestamp if splitting mid key is enabled.
@@ -157,6 +174,7 @@ func evalExport(
 			MaxSize:            maxSize,
 			StopMidKey:         args.SplitMidKey,
 			UseTBI:             useTBI,
+			ResourceLimiter:    storage.NewResourceLimiter(storage.ResourceLimiterOptions{MaxRunTime: maxRunTime}, timeutil.DefaultTimeSource{}),
 		}, destFile)
 		if err != nil {
 			if errors.HasType(err, (*storage.ExceedMaxSizeError)(nil)) {