kvserver: add scheduler re-add test

Author: tbg

pkg/kv/kvserver/scheduler.go

@@ -14,6 +14,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvflowcontrol/rac2"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/util"
+	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
@@ -22,6 +23,12 @@ import (
 
 const rangeIDChunkSize = 1000
 
+type testProcessorI interface {
+	processTestEvent(
+		context.Context, roachpb.RangeID, *raftSchedulerShard, raftScheduleState,
+	)
+}
+
 type rangeIDChunk struct {
 	// Valid contents are buf[rd:wr], read at buf[rd], write at buf[wr].
 	buf    [rangeIDChunkSize]roachpb.RangeID
@@ -134,6 +141,7 @@ const (
 	stateRaftTick
 	stateRACv2PiggybackedAdmitted
 	stateRACv2RangeController
+	stateTestIntercept // used for testing, CrdbTestBuild only
 )
 
 type raftScheduleState struct {
@@ -424,6 +432,9 @@ func (ss *raftSchedulerShard) worker(
 		if state.flags&stateRACv2RangeController != 0 {
 			processor.processRACv2RangeController(ctx, id)
 		}
+		if buildutil.CrdbTestBuild && state.flags&stateTestIntercept != 0 {
+			processor.(testProcessorI).processTestEvent(ctx, id, ss, state)
+		}
 
 		ss.Lock()
 		state = ss.state[id]

pkg/kv/kvserver/scheduler_test.go

@@ -26,6 +26,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 	"github.com/cockroachdb/errors"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -135,15 +136,19 @@ type testProcessor struct {
 		rac2RangeController     map[roachpb.RangeID]int
 		ready                   func(roachpb.RangeID)
 	}
+	testEventCh chan func(roachpb.RangeID, *raftSchedulerShard, raftScheduleState)
 }
 
+var _ testProcessorI = (*testProcessor)(nil)
+
 func newTestProcessor() *testProcessor {
 	p := &testProcessor{}
 	p.mu.raftReady = make(map[roachpb.RangeID]int)
 	p.mu.raftRequest = make(map[roachpb.RangeID]int)
 	p.mu.raftTick = make(map[roachpb.RangeID]int)
 	p.mu.rac2PiggybackedAdmitted = make(map[roachpb.RangeID]int)
 	p.mu.rac2RangeController = make(map[roachpb.RangeID]int)
+	p.testEventCh = make(chan func(roachpb.RangeID, *raftSchedulerShard, raftScheduleState), 10)
 	return p
 }
 
@@ -192,6 +197,16 @@ func (p *testProcessor) processRACv2RangeController(_ context.Context, rangeID r
 	p.mu.Unlock()
 }
 
+func (p *testProcessor) processTestEvent(
+	ctx context.Context, id roachpb.RangeID, ss *raftSchedulerShard, ev raftScheduleState,
+) {
+	select {
+	case fn := <-p.testEventCh:
+		fn(id, ss, ev)
+	default:
+	}
+}
+
 func (p *testProcessor) readyCount(rangeID roachpb.RangeID) int {
 	p.mu.Lock()
 	defer p.mu.Unlock()
@@ -352,6 +367,75 @@ func TestSchedulerBuffering(t *testing.T) {
 	}
 }
 
+func TestSchedulerEnqueueWhileProcessing(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	skip.UnderNonTestBuild(t) // stateTestIntercept needs crdb test build
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	stopper := stop.NewStopper()
+	defer stopper.Stop(ctx)
+
+	m := newStoreMetrics(metric.TestSampleInterval)
+	p := newTestProcessor()
+	s := newRaftScheduler(log.MakeTestingAmbientContext(stopper.Tracer()), m, p, 1, 1, 1, 5)
+	s.Start(stopper)
+
+	done := make(chan struct{})
+
+	// Inject code into the "middle" of event processing - after having consumed
+	// from the queue, but before re-checking of overlapping enqueue calls.
+	p.testEventCh <- func(id roachpb.RangeID, ss *raftSchedulerShard, ev raftScheduleState) {
+		// First call into this method.
+		//
+		// The event calling into us must have `ev.begin` set; it was set when
+		// enqueuing.
+		assert.NotZero(t, ev.begin)
+
+		// Even though our event is currently being processed, there is a queued
+		// and otherwise blank event in the scheduler state (which is how we have
+		// concurrent enqueue calls coalesce onto the still pending processing of
+		// the current event).
+		ss.Lock()
+		statePre := ss.state[id]
+		ss.Unlock()
+
+		assert.Zero(t, statePre.begin)
+		assert.Equal(t, stateQueued, statePre.flags)
+
+		// Simulate a concurrent actor that enqueues the same range again.
+		// This will not trigger the interceptor again, since the done channel
+		// is closed by that time.
+		s.enqueue1(stateTestIntercept, 1)
+
+		// Seeing that there is an existing "queued" event, the enqueue call below
+		// should not populate `begin`.  Instead, this will be the job of our
+		// caller when it *actually* pushes into the queue again after fully
+		// having handled `ev`.
+		ss.Lock()
+		statePost := ss.state[id]
+		ss.Unlock()
+
+		// TODO(tbg): enable in follow-up commit.
+		// assert.Zero(t, statePost.begin)
+		assert.Equal(t, stateQueued|stateTestIntercept, statePost.flags)
+		close(done)
+	}
+	p.testEventCh <- func(id roachpb.RangeID, shard *raftSchedulerShard, ev raftScheduleState) {
+		// Second call into this method, i.e. the overlappingly-enqeued event is
+		// being processed. Check that `begin` is now set.
+		assert.NotZero(t, ev.begin)
+	}
+	s.enqueue1(stateTestIntercept, 1) // will become 'ev' in the intercept
+	select {
+	case <-done:
+	case <-ctx.Done():
+		t.Fatal(ctx.Err())
+	}
+}
+
 func TestNewSchedulerShards(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)