Skip to content

Commit

Permalink
Cherry-pick ad7bdd9 with conflicts
Browse files Browse the repository at this point in the history
  • Loading branch information
vitess-bot[bot] committed Mar 12, 2024
1 parent d3cd5d2 commit 87925bb
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 4 deletions.
13 changes: 9 additions & 4 deletions go/test/endtoend/vtgate/queries/union/union_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@ import (
"testing"

"vitess.io/vitess/go/test/endtoend/cluster"

"vitess.io/vitess/go/test/endtoend/utils"

"github.com/stretchr/testify/assert"
Expand Down Expand Up @@ -103,9 +102,15 @@ func TestUnionAll(t *testing.T) {
mcmp.AssertMatchesNoOrder("select tbl2.id1 FROM ((select id1 from t1 order by id1 limit 5) union all (select id1 from t1 order by id1 desc limit 5)) as tbl1 INNER JOIN t1 as tbl2 ON tbl1.id1 = tbl2.id1",
"[[INT64(1)] [INT64(2)] [INT64(2)] [INT64(1)]]")

// union all between two select unique in tables
mcmp.AssertMatchesNoOrder("select id1 from t1 where id1 in (1, 2, 3, 4, 5, 6, 7, 8) union all select id1 from t1 where id1 in (1, 2, 3, 4, 5, 6, 7, 8)",
"[[INT64(1)] [INT64(2)] [INT64(1)] [INT64(2)]]")
// this test is quite good at uncovering races in the Concatenate engine primitive. make it run many times
// see: https://github.com/vitessio/vitess/issues/15434
if utils.BinaryIsAtLeastAtVersion(20, "vtgate") {
for i := 0; i < 100; i++ {
// union all between two select unique in tables
mcmp.AssertMatchesNoOrder("select id1 from t1 where id1 in (1, 2, 3, 4, 5, 6, 7, 8) union all select id1 from t1 where id1 in (1, 2, 3, 4, 5, 6, 7, 8)",
"[[INT64(1)] [INT64(2)] [INT64(1)] [INT64(2)]]")
}
}

// 4 tables union all
mcmp.AssertMatchesNoOrder("select id1, id2 from t1 where id1 = 1 union all select id3,id4 from t2 where id3 = 3 union all select id1, id2 from t1 where id1 = 2 union all select id3,id4 from t2 where id3 = 4",
Expand Down
36 changes: 36 additions & 0 deletions go/vt/vtgate/engine/concatenate.go
Original file line number Diff line number Diff line change
Expand Up @@ -215,6 +215,7 @@ func (c *Concatenate) parallelStreamExec(ctx context.Context, vcursor VCursor, b
for i, source := range c.Sources {
wg.Add(1)
currIndex, currSource := i, source
<<<<<<< HEAD

go func() {
defer wg.Done()
Expand Down Expand Up @@ -244,6 +245,41 @@ func (c *Concatenate) parallelStreamExec(ctx context.Context, vcursor VCursor, b
defer cbMu.Unlock()
select {
case <-ctx.Done():
=======
wg.Go(func() error {
err := vcursor.StreamExecutePrimitive(ctx, currSource, bindVars, true, func(resultChunk *sqltypes.Result) error {
muFields.Lock()

// Process fields when they arrive; coordinate field agreement across sources.
if resultChunk.Fields != nil && rest[currIndex] == nil {
// Capture the initial result chunk to determine field types later.
rest[currIndex] = resultChunk

// If this was the last source to report its fields, derive the final output fields.
if !slices.Contains(rest, nil) {
// We have received fields from all sources. We can now calculate the output types
var err error
resultChunk.Fields, fieldTypes, err = c.getFieldTypes(vcursor, rest)
if err != nil {
muFields.Unlock()
return err
}

muFields.Unlock()
defer condFields.Broadcast()
return callback(resultChunk, currIndex)
}
}

// Wait for fields from all sources.
for slices.Contains(rest, nil) {
condFields.Wait()
}
muFields.Unlock()

// Context check to avoid extra work.
if ctx.Err() != nil {
>>>>>>> ad7bdd9380 (engine: fix race in concatenate (#15454))
return nil
default:
return callback(resultChunk)
Expand Down

0 comments on commit 87925bb

Please sign in to comment.