func (ds *ServerImpl) setupFlow(
ctx context.Context, req *SetupFlowRequest, simpleFlowConsumer RowReceiver,
) (*Flow, error) {
sp, err := tracing.JoinOrNew(ds.AmbientContext.Tracer, req.TraceContext, "flow")
if err != nil {
return nil, err
}
ctx = opentracing.ContextWithSpan(ctx, sp)
txn := ds.setupTxn(ctx, &req.Txn)
flowCtx := FlowCtx{
Context: ctx,
id: req.Flow.FlowID,
evalCtx: &ds.evalCtx,
rpcCtx: ds.RPCContext,
txn: txn,
}
f := newFlow(flowCtx, ds.flowRegistry, simpleFlowConsumer)
if err := f.setupFlow(&req.Flow); err != nil {
log.Error(ctx, err)
sp.Finish()
return nil, err
}
return f, nil
}
func (ds *ServerImpl) setupFlow(
ctx context.Context, req *SetupFlowRequest, syncFlowConsumer RowReceiver,
) (*Flow, error) {
sp, err := tracing.JoinOrNew(ds.AmbientContext.Tracer, req.TraceContext, "flow")
if err != nil {
return nil, err
}
ctx = opentracing.ContextWithSpan(ctx, sp)
// TODO(radu): we should sanity check some of these fields (especially
// txnProto).
flowCtx := FlowCtx{
Context: ctx,
id: req.Flow.FlowID,
evalCtx: &ds.evalCtx,
rpcCtx: ds.RPCContext,
txnProto: &req.Txn,
clientDB: ds.DB,
}
f := newFlow(flowCtx, ds.flowRegistry, syncFlowConsumer)
if err := f.setupFlow(&req.Flow); err != nil {
log.Error(ctx, err)
sp.Finish()
return nil, err
}
return f, nil
}
func (n *Node) batchInternal(
ctx context.Context, args *roachpb.BatchRequest,
) (*roachpb.BatchResponse, error) {
// TODO(marc): grpc's authentication model (which gives credential access in
// the request handler) doesn't really fit with the current design of the
// security package (which assumes that TLS state is only given at connection
// time) - that should be fixed.
if peer, ok := peer.FromContext(ctx); ok {
if tlsInfo, ok := peer.AuthInfo.(credentials.TLSInfo); ok {
certUser, err := security.GetCertificateUser(&tlsInfo.State)
if err != nil {
return nil, err
}
if certUser != security.NodeUser {
return nil, errors.Errorf("user %s is not allowed", certUser)
}
}
}
var br *roachpb.BatchResponse
type snowballInfo struct {
syncutil.Mutex
collectedSpans [][]byte
done bool
}
var snowball *snowballInfo
if err := n.stopper.RunTaskWithErr(func() error {
const opName = "node.Batch"
sp, err := tracing.JoinOrNew(n.storeCfg.AmbientCtx.Tracer, args.TraceContext, opName)
if err != nil {
return err
}
// If this is a snowball span, it gets special treatment: It skips the
// regular tracing machinery, and we instead send the collected spans
// back with the response. This is more expensive, but then again,
// those are individual requests traced by users, so they can be.
if sp.BaggageItem(tracing.Snowball) != "" {
sp.LogEvent("delegating to snowball tracing")
sp.Finish()
snowball = new(snowballInfo)
recorder := func(rawSpan basictracer.RawSpan) {
snowball.Lock()
defer snowball.Unlock()
if snowball.done {
// This is a late span that we must discard because the request was
// already completed.
return
}
encSp, err := tracing.EncodeRawSpan(&rawSpan, nil)
if err != nil {
log.Warning(ctx, err)
}
snowball.collectedSpans = append(snowball.collectedSpans, encSp)
}
if sp, err = tracing.JoinOrNewSnowball(opName, args.TraceContext, recorder); err != nil {
return err
}
}
defer sp.Finish()
traceCtx := opentracing.ContextWithSpan(ctx, sp)
log.Event(traceCtx, args.Summary())
tStart := timeutil.Now()
var pErr *roachpb.Error
br, pErr = n.stores.Send(traceCtx, *args)
if pErr != nil {
br = &roachpb.BatchResponse{}
log.ErrEventf(traceCtx, "%T", pErr.GetDetail())
}
if br.Error != nil {
panic(roachpb.ErrorUnexpectedlySet(n.stores, br))
}
n.metrics.callComplete(timeutil.Since(tStart), pErr)
br.Error = pErr
return nil
}); err != nil {
return nil, err
}
if snowball != nil {
snowball.Lock()
br.CollectedSpans = snowball.collectedSpans
snowball.done = true
snowball.Unlock()
}
return br, nil
}
func TestClusterFlow(t *testing.T) {
defer leaktest.AfterTest(t)()
const numRows = 100
args := base.TestClusterArgs{ReplicationMode: base.ReplicationManual}
tc := serverutils.StartTestCluster(t, 3, args)
defer tc.Stopper().Stop()
sumDigitsFn := func(row int) parser.Datum {
sum := 0
for row > 0 {
sum += row % 10
row /= 10
}
return parser.NewDInt(parser.DInt(sum))
}
sqlutils.CreateTable(t, tc.ServerConn(0), "t",
"num INT PRIMARY KEY, digitsum INT, numstr STRING, INDEX s (digitsum)",
numRows,
sqlutils.ToRowFn(sqlutils.RowIdxFn, sumDigitsFn, sqlutils.RowEnglishFn))
kvDB := tc.Server(0).KVClient().(*client.DB)
desc := sqlbase.GetTableDescriptor(kvDB, "test", "t")
makeIndexSpan := func(start, end int) TableReaderSpan {
var span roachpb.Span
prefix := roachpb.Key(sqlbase.MakeIndexKeyPrefix(desc, desc.Indexes[0].ID))
span.Key = append(prefix, encoding.EncodeVarintAscending(nil, int64(start))...)
span.EndKey = append(span.EndKey, prefix...)
span.EndKey = append(span.EndKey, encoding.EncodeVarintAscending(nil, int64(end))...)
return TableReaderSpan{Span: span}
}
// Set up table readers on three hosts feeding data into a join reader on
// the third host. This is a basic test for the distributed flow
// infrastructure, including local and remote streams.
//
// Note that the ranges won't necessarily be local to the table readers, but
// that doesn't matter for the purposes of this test.
// Start a span (useful to look at spans using Lighstep).
sp, err := tracing.JoinOrNew(tracing.NewTracer(), nil, "cluster test")
if err != nil {
t.Fatal(err)
}
ctx := opentracing.ContextWithSpan(context.Background(), sp)
defer sp.Finish()
tr1 := TableReaderSpec{
Table: *desc,
IndexIdx: 1,
OutputColumns: []uint32{0, 1},
Spans: []TableReaderSpan{makeIndexSpan(0, 8)},
}
tr2 := TableReaderSpec{
Table: *desc,
IndexIdx: 1,
OutputColumns: []uint32{0, 1},
Spans: []TableReaderSpan{makeIndexSpan(8, 12)},
}
tr3 := TableReaderSpec{
Table: *desc,
IndexIdx: 1,
OutputColumns: []uint32{0, 1},
Spans: []TableReaderSpan{makeIndexSpan(12, 100)},
}
jr := JoinReaderSpec{
Table: *desc,
OutputColumns: []uint32{2},
}
txn := client.NewTxn(ctx, *kvDB)
fid := FlowID{uuid.MakeV4()}
req1 := &SetupFlowRequest{Txn: txn.Proto}
req1.Flow = FlowSpec{
FlowID: fid,
Processors: []ProcessorSpec{{
Core: ProcessorCoreUnion{TableReader: &tr1},
Output: []OutputRouterSpec{{
Type: OutputRouterSpec_MIRROR,
Streams: []StreamEndpointSpec{
{StreamID: 0, Mailbox: &MailboxSpec{TargetAddr: tc.Server(2).ServingAddr()}},
},
}},
}},
}
req2 := &SetupFlowRequest{Txn: txn.Proto}
req2.Flow = FlowSpec{
FlowID: fid,
Processors: []ProcessorSpec{{
Core: ProcessorCoreUnion{TableReader: &tr2},
Output: []OutputRouterSpec{{
Type: OutputRouterSpec_MIRROR,
Streams: []StreamEndpointSpec{
{StreamID: 1, Mailbox: &MailboxSpec{TargetAddr: tc.Server(2).ServingAddr()}},
},
}},
}},
}
req3 := &SetupFlowRequest{Txn: txn.Proto}
req3.Flow = FlowSpec{
FlowID: fid,
Processors: []ProcessorSpec{
{
Core: ProcessorCoreUnion{TableReader: &tr3},
Output: []OutputRouterSpec{{
Type: OutputRouterSpec_MIRROR,
Streams: []StreamEndpointSpec{
{StreamID: StreamID(2)},
},
}},
},
{
Input: []InputSyncSpec{{
Type: InputSyncSpec_ORDERED,
Ordering: Ordering{Columns: []Ordering_Column{{1, Ordering_Column_ASC}}},
Streams: []StreamEndpointSpec{
{StreamID: 0, Mailbox: &MailboxSpec{}},
{StreamID: 1, Mailbox: &MailboxSpec{}},
{StreamID: StreamID(2)},
},
}},
Core: ProcessorCoreUnion{JoinReader: &jr},
Output: []OutputRouterSpec{{
Type: OutputRouterSpec_MIRROR,
Streams: []StreamEndpointSpec{{Mailbox: &MailboxSpec{SimpleResponse: true}}},
}}},
},
}
if err := SetFlowRequestTrace(ctx, req1); err != nil {
t.Fatal(err)
}
if err := SetFlowRequestTrace(ctx, req2); err != nil {
t.Fatal(err)
}
if err := SetFlowRequestTrace(ctx, req3); err != nil {
t.Fatal(err)
}
var clients []DistSQLClient
for i := 0; i < 3; i++ {
s := tc.Server(i)
conn, err := s.RPCContext().GRPCDial(s.ServingAddr())
if err != nil {
t.Fatal(err)
}
clients = append(clients, NewDistSQLClient(conn))
}
if log.V(1) {
log.Infof(ctx, "Setting up flow on 0")
}
if resp, err := clients[0].SetupFlow(ctx, req1); err != nil {
t.Fatal(err)
} else if resp.Error != nil {
t.Fatal(resp.Error)
}
if log.V(1) {
log.Infof(ctx, "Setting up flow on 1")
}
if resp, err := clients[1].SetupFlow(ctx, req2); err != nil {
t.Fatal(err)
} else if resp.Error != nil {
t.Fatal(resp.Error)
}
if log.V(1) {
log.Infof(ctx, "Running flow on 2")
}
stream, err := clients[2].RunSimpleFlow(ctx, req3)
if err != nil {
t.Fatal(err)
}
var decoder StreamDecoder
var rows sqlbase.EncDatumRows
for {
msg, err := stream.Recv()
if err != nil {
if err == io.EOF {
break
}
t.Fatal(err)
}
err = decoder.AddMessage(msg)
if err != nil {
t.Fatal(err)
}
rows = testGetDecodedRows(t, &decoder, rows)
}
if done, trailerErr := decoder.IsDone(); !done {
t.Fatal("stream not done")
} else if trailerErr != nil {
t.Fatal("error in the stream trailer:", trailerErr)
}
// The result should be all the numbers in string form, ordered by the
// digit sum (and then by number).
var results []string
for sum := 1; sum <= 50; sum++ {
for i := 1; i <= numRows; i++ {
if int(*sumDigitsFn(i).(*parser.DInt)) == sum {
results = append(results, fmt.Sprintf("['%s']", sqlutils.IntToEnglish(i)))
}
}
}
expected := strings.Join(results, " ")
expected = "[" + expected + "]"
if rowStr := rows.String(); rowStr != expected {
t.Errorf("Result: %s\n Expected: %s\n", rowStr, expected)
}
}
