storage: err on snapshot of large ranges

If a range needs to be split, return an err rather than attempting to generate a snapshot.
This avoids generating excessively large snapshots.

Suggested in cockroachdb#7581.

storage/client_bench_test.go

@@ -14,50 +14,29 @@
 //
 // Author: Tamir Duberstein ([email protected])
 
-package storage_test
+package storage
 
 import (
-	"math/rand"
 	"testing"
 
-	"github.com/cockroachdb/cockroach/internal/client"
-	"github.com/cockroachdb/cockroach/keys"
-	"github.com/cockroachdb/cockroach/roachpb"
-	"github.com/cockroachdb/cockroach/storage"
-	"github.com/cockroachdb/cockroach/util/randutil"
 	"github.com/cockroachdb/cockroach/util/tracing"
 )
 
 func BenchmarkReplicaSnapshot(b *testing.B) {
 	defer tracing.Disable()()
-	sCtx := storage.TestStoreContext()
+	sCtx := TestStoreContext()
 	sCtx.TestingKnobs.DisableSplitQueue = true
-	store, stopper, _ := createTestStoreWithContext(b, sCtx)
+	store, _, stopper := createTestStoreWithContext(b, &sCtx)
 	defer stopper.Stop()
 	// We want to manually control the size of the raft log.
 	store.SetRaftLogQueueActive(false)
 
-	const rangeID = 1
-	const keySize = 1 << 7   // 128 B
-	const valSize = 1 << 10  // 1 KiB
-	const snapSize = 1 << 25 // 32 MiB
-
 	rep, err := store.GetReplica(rangeID)
 	if err != nil {
 		b.Fatal(err)
 	}
 
-	src := rand.New(rand.NewSource(0))
-	for i := 0; i < snapSize/(keySize+valSize); i++ {
-		key := keys.MakeRowSentinelKey(randutil.RandBytes(src, keySize))
-		val := randutil.RandBytes(src, valSize)
-		pArgs := putArgs(key, val)
-		if _, pErr := client.SendWrappedWith(rep, nil, roachpb.Header{
-			RangeID: rangeID,
-		}, &pArgs); pErr != nil {
-			b.Fatal(pErr)
-		}
-	}
+	fillTestRange(b, rep, snapSize)
 
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {

storage/replica.go

@@ -2612,6 +2612,12 @@ func (r *Replica) needsSplitBySize() bool {
 	return maxBytes > 0 && size > maxBytes
 }
 
+func (r *Replica) exceedsDoubleSplitSizeLocked() bool {
+	maxBytes := r.mu.maxBytes
+	size := r.mu.state.Stats.Total()
+	return maxBytes > 0 && size > maxBytes*2
+}
+
 // maybeAddToSplitQueue checks whether the current size of the range
 // exceeds the max size specified in the zone config. If yes, the
 // range is added to the split queue.

storage/replica_raftstorage.go

@@ -247,6 +247,10 @@ func (r *Replica) GetFirstIndex() (uint64, error) {
 // Snapshot implements the raft.Storage interface.
 // Snapshot requires that the replica lock is held.
 func (r *Replica) Snapshot() (raftpb.Snapshot, error) {
+	if r.exceedsDoubleSplitSizeLocked() {
+		return raftpb.Snapshot{}, raft.ErrSnapshotTemporarilyUnavailable
+	}
+
 	rangeID := r.RangeID
 
 	// If a snapshot is in progress, see if it's ready.

storage/replica_raftstorage_test.go

@@ -17,13 +17,18 @@
 package storage
 
 import (
+	"math/rand"
 	"testing"
 
 	"golang.org/x/net/context"
 
+	"github.com/cockroachdb/cockroach/internal/client"
+	"github.com/cockroachdb/cockroach/keys"
 	"github.com/cockroachdb/cockroach/roachpb"
 	"github.com/cockroachdb/cockroach/testutils"
 	"github.com/cockroachdb/cockroach/util/leaktest"
+	"github.com/cockroachdb/cockroach/util/randutil"
+	"github.com/coreos/etcd/raft"
 	"github.com/coreos/etcd/raft/raftpb"
 )
 
@@ -81,3 +86,51 @@ func TestApplySnapshotDenyPreemptive(t *testing.T) {
 	}
 
 }
+
+const rangeID = 1
+const keySize = 1 << 7   // 128 B
+const valSize = 1 << 10  // 1 KiB
+const snapSize = 1 << 25 // 32 MiB
+
+func fillTestRange(t testing.TB, rep *Replica, size int) {
+	src := rand.New(rand.NewSource(0))
+	for i := 0; i < snapSize/(keySize+valSize); i++ {
+		key := keys.MakeRowSentinelKey(randutil.RandBytes(src, keySize))
+		val := randutil.RandBytes(src, valSize)
+		pArgs := putArgs(key, val)
+		if _, pErr := client.SendWrappedWith(rep, nil, roachpb.Header{
+			RangeID: rangeID,
+		}, &pArgs); pErr != nil {
+			t.Fatal(pErr)
+		}
+	}
+}
+
+func TestSkipLargeReplicaSnapshot(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	store, _, stopper := createTestStore(t)
+	store.ctx.TestingKnobs.DisableSplitQueue = true
+	// We want to manually control the size of the raft log.
+	defer stopper.Stop()
+
+	rep, err := store.GetReplica(rangeID)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if pErr := rep.redirectOnOrAcquireLease(context.Background()); pErr != nil {
+		t.Fatal(pErr)
+	}
+
+	fillTestRange(t, rep, snapSize)
+
+	if _, err := rep.Snapshot(); err != nil {
+		t.Fatal(err)
+	}
+
+	fillTestRange(t, rep, snapSize*2)
+
+	if _, err := rep.Snapshot(); err != raft.ErrSnapshotTemporarilyUnavailable {
+		t.Fatalf("snapshot of a very large range should fail but got %v", err)
+	}
+}