github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/distsql/sync_flow_after_drain_test.go (about)

     1  // Copyright 2019 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package distsql
    12  
    13  import (
    14  	"context"
    15  	"testing"
    16  	"time"
    17  
    18  	"github.com/cockroachdb/cockroach/pkg/base"
    19  	"github.com/cockroachdb/cockroach/pkg/sql/execinfra"
    20  	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
    21  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    22  	"github.com/cockroachdb/cockroach/pkg/testutils"
    23  	"github.com/cockroachdb/cockroach/pkg/testutils/distsqlutils"
    24  	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
    25  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    26  )
    27  
    28  // Test that we can register send a sync flow to the distSQLSrv after the
    29  // FlowRegistry is draining and the we can also clean that flow up (the flow
    30  // will get a draining error). This used to crash.
    31  func TestSyncFlowAfterDrain(t *testing.T) {
    32  	defer leaktest.AfterTest(t)()
    33  
    34  	ctx := context.Background()
    35  	// We'll create a server just so that we can extract its distsql ServerConfig,
    36  	// so we can use it for a manually-built DistSQL Server below. Otherwise, too
    37  	// much work to create that ServerConfig by hand.
    38  	s, _, _ := serverutils.StartServer(t, base.TestServerArgs{})
    39  	defer s.Stopper().Stop(ctx)
    40  	cfg := s.DistSQLServer().(*ServerImpl).ServerConfig
    41  
    42  	distSQLSrv := NewServer(ctx, cfg)
    43  	distSQLSrv.flowRegistry.Drain(
    44  		time.Duration(0) /* flowDrainWait */, time.Duration(0) /* minFlowDrainWait */, nil /* reporter */)
    45  
    46  	// We create some flow; it doesn't matter what.
    47  	req := execinfrapb.SetupFlowRequest{Version: execinfra.Version}
    48  	req.Flow = execinfrapb.FlowSpec{
    49  		Processors: []execinfrapb.ProcessorSpec{
    50  			{
    51  				Core: execinfrapb.ProcessorCoreUnion{Values: &execinfrapb.ValuesCoreSpec{}},
    52  				Output: []execinfrapb.OutputRouterSpec{{
    53  					Type:    execinfrapb.OutputRouterSpec_PASS_THROUGH,
    54  					Streams: []execinfrapb.StreamEndpointSpec{{StreamID: 1, Type: execinfrapb.StreamEndpointSpec_REMOTE}},
    55  				}},
    56  			},
    57  			{
    58  				Input: []execinfrapb.InputSyncSpec{{
    59  					Type:    execinfrapb.InputSyncSpec_UNORDERED,
    60  					Streams: []execinfrapb.StreamEndpointSpec{{StreamID: 1, Type: execinfrapb.StreamEndpointSpec_REMOTE}},
    61  				}},
    62  				Core: execinfrapb.ProcessorCoreUnion{Noop: &execinfrapb.NoopCoreSpec{}},
    63  				Output: []execinfrapb.OutputRouterSpec{{
    64  					Type:    execinfrapb.OutputRouterSpec_PASS_THROUGH,
    65  					Streams: []execinfrapb.StreamEndpointSpec{{Type: execinfrapb.StreamEndpointSpec_SYNC_RESPONSE}},
    66  				}},
    67  			},
    68  		},
    69  	}
    70  
    71  	types := make([]*types.T, 0)
    72  	rb := distsqlutils.NewRowBuffer(types, nil /* rows */, distsqlutils.RowBufferArgs{})
    73  	ctx, flow, err := distSQLSrv.SetupSyncFlow(ctx, &distSQLSrv.memMonitor, &req, rb)
    74  	if err != nil {
    75  		t.Fatal(err)
    76  	}
    77  	if err := flow.Start(ctx, func() {}); err != nil {
    78  		t.Fatal(err)
    79  	}
    80  	flow.Wait()
    81  	_, meta := rb.Next()
    82  	if meta == nil {
    83  		t.Fatal("expected draining err, got no meta")
    84  	}
    85  	if !testutils.IsError(meta.Err, "the registry is draining") {
    86  		t.Fatalf("expected draining err, got: %v", meta.Err)
    87  	}
    88  	flow.Cleanup(ctx)
    89  }