Load TPC-H concurrently

ch/workload.go

@@ -107,9 +107,9 @@ func (w Workloader) Prepare(ctx context.Context, threadID int) error {
 		return err
 	}
 	sqlLoader := map[dbgen.Table]dbgen.Loader{
-		dbgen.TSupp:   tpch.NewSuppLoader(ctx, w.db),
-		dbgen.TNation: tpch.NewNationLoader(ctx, w.db),
-		dbgen.TRegion: tpch.NewRegionLoader(ctx, w.db),
+		dbgen.TSupp:   tpch.NewSuppLoader(ctx, w.db, 1),
+		dbgen.TNation: tpch.NewNationLoader(ctx, w.db, 1),
+		dbgen.TRegion: tpch.NewRegionLoader(ctx, w.db, 1),
 	}
 	dbgen.InitDbGen(1)
 	if err := dbgen.DbGen(sqlLoader, []dbgen.Table{dbgen.TNation, dbgen.TRegion, dbgen.TSupp}); err != nil {

cmd/go-tpc/tpch.go

@@ -6,8 +6,9 @@ import (
 	"os"
 	"strings"
 
-	"github.com/pingcap/go-tpc/tpch"
 	"github.com/spf13/cobra"
+
+	"github.com/pingcap/go-tpc/tpch"
 )
 
 var tpchConfig tpch.Config
@@ -16,13 +17,13 @@ func executeTpch(action string) {
 	openDB()
 	defer closeDB()
 
-	// if globalDB == nil
 	if globalDB == nil {
 		fmt.Fprintln(os.Stderr, "cannot connect to the database")
 		os.Exit(1)
 	}
 
 	tpchConfig.DBName = dbName
+	tpchConfig.PrepareThreads = threads
 	tpchConfig.QueryNames = strings.Split(tpchConfig.RawQueries, ",")
 	w := tpch.NewWorkloader(globalDB, &tpchConfig)
 

go.mod

@@ -10,4 +10,7 @@ require (
 	github.com/mattn/go-runewidth v0.0.7 // indirect
 	github.com/prometheus/client_golang v0.9.3
 	github.com/spf13/cobra v1.0.0
+	github.com/stretchr/testify v1.6.1
+	go.uber.org/atomic v1.9.0
+	golang.org/x/sync v0.0.0-20210220032951-036812b2e83c
 )

go.sum

@@ -380,6 +380,8 @@ go.mongodb.org/mongo-driver v1.4.3/go.mod h1:WcMNYLx/IlOxLe6JRJiv2uXuCz6zBLndR4S
 go.opencensus.io v0.21.0/go.mod h1:mSImk1erAIZhrmZN+AvHh14ztQfjbGwt4TtuofqLduU=
 go.opencensus.io v0.22.0/go.mod h1:+kGneAE2xo2IficOXnaByMWTGM9T73dGwxeWcUqIpI8=
 go.uber.org/atomic v1.4.0/go.mod h1:gD2HeocX3+yG+ygLZcrzQJaqmWj9AIm7n08wl/qW/PE=
+go.uber.org/atomic v1.9.0 h1:ECmE8Bn/WFTYwEW/bpKD3M8VtR/zQVbavAoalC1PYyE=
+go.uber.org/atomic v1.9.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
 go.uber.org/multierr v1.1.0/go.mod h1:wR5kodmAFQ0UK8QlbwjlSNy0Z68gJhDJUG5sjR94q/0=
 go.uber.org/zap v1.10.0/go.mod h1:vwi/ZaCAaUcBkycHslxD9B2zi4UTXhF60s6SWpuDF0Q=
 golang.org/x/crypto v0.0.0-20180904163835-0709b304e793/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4=
@@ -444,6 +446,8 @@ golang.org/x/sync v0.0.0-20190412183630-56d357773e84/go.mod h1:RxMgew5VJxzue5/jJ
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20200625203802-6e8e738ad208/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.0.0-20210220032951-036812b2e83c h1:5KslGYwFpkhGh+Q16bwMP3cOontH8FOep7tGV86Y7SQ=
+golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20180823144017-11551d06cbcc/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=

pkg/sink/concurrent.go

@@ -0,0 +1,139 @@
+package sink
+
+import (
+	"context"
+	"sync"
+
+	"go.uber.org/atomic"
+	"golang.org/x/sync/errgroup"
+)
+
+// ConcurrentSink inserts values to one of the downstream sinks.
+// The insert will be blocked if all downstream sinks are working.
+//
+// WARN: Although this sink can transform serial Writes to multiple down stream sinks, this sink itself is not
+// concurrent safe. You must not call WriteRow and Flush concurrently.
+type ConcurrentSink struct {
+	allSinks []Sink
+
+	writeCh       chan writeRowOp
+	writeResultCh chan error
+	writeWg       sync.WaitGroup
+
+	concurrentGuard atomic.Int32 // Used to check whether this struct is used concurrently
+}
+
+type writeRowOp struct {
+	ctx    context.Context
+	values []interface{}
+}
+
+var _ Sink = &ConcurrentSink{}
+
+func NewConcurrentSink(downStreamBuilder func(idx int) Sink, concurrency int) *ConcurrentSink {
+	sinks := make([]Sink, concurrency)
+	for i := 0; i < concurrency; i++ {
+		sinks[i] = downStreamBuilder(i)
+	}
+
+	cs := &ConcurrentSink{
+		allSinks:      sinks,
+		writeCh:       make(chan writeRowOp, concurrency),
+		writeResultCh: make(chan error, 1),
+	}
+	for i := 0; i < concurrency; i++ {
+		go cs.runConsumerLoop(i)
+	}
+	return cs
+}
+
+func (c *ConcurrentSink) runConsumerLoop(downStreamIdx int) {
+	sink := c.allSinks[downStreamIdx]
+
+	for {
+		select {
+		case op, ok := <-c.writeCh:
+			if !ok {
+				// Channel close
+				return
+			}
+			err := sink.WriteRow(op.ctx, op.values...)
+			c.writeWg.Add(-1)
+			if err != nil {
+				select {
+				case c.writeResultCh <- err:
+				default:
+				}
+			}
+		}
+	}
+}
+
+func (c *ConcurrentSink) WriteRow(ctx context.Context, values ...interface{}) error {
+	v := c.concurrentGuard.Inc()
+	if v > 1 {
+		panic("ConcurrentSink cannot be called concurrently")
+	}
+	defer c.concurrentGuard.Dec()
+
+	c.writeWg.Add(1)
+	c.writeCh <- writeRowOp{
+		ctx:    ctx,
+		values: values,
+	}
+	select {
+	case err := <-c.writeResultCh:
+		return err
+	default:
+		return nil
+	}
+}
+
+// Flush flushes all downstream sinks concurrently, wait all sinks to be flushed and returns the first error
+// encountered.
+//
+// WARN: Flush() will wait until all existing write ops are finished.
+func (c *ConcurrentSink) Flush(ctx context.Context) error {
+	v := c.concurrentGuard.Inc()
+	if v > 1 {
+		panic("ConcurrentSink cannot be called concurrently")
+	}
+	defer c.concurrentGuard.Dec()
+
+	// Wait all writes to finish.
+	c.writeWg.Wait()
+
+	// At this time there is no running write ops, so we are safe to call sink.Flush() for each sink.
+	g, ctx := errgroup.WithContext(ctx)
+	for _, sink_ := range c.allSinks {
+		sink := sink_
+		g.Go(func() error {
+			return sink.Flush(ctx)
+		})
+	}
+	return g.Wait()
+}
+
+// Close closes all downstream sinks concurrently, wait all sinks to be closed and returns the first error
+// encountered.
+//
+// WARN: Close() will wait until all existing write ops are finished.
+func (c *ConcurrentSink) Close(ctx context.Context) error {
+	v := c.concurrentGuard.Inc()
+	if v > 1 {
+		panic("ConcurrentSink cannot be called concurrently")
+	}
+	defer c.concurrentGuard.Dec()
+
+	// Wait all writes to finish.
+	c.writeWg.Wait()
+
+	g := new(errgroup.Group)
+	for _, sink_ := range c.allSinks {
+		sink := sink_
+		g.Go(func() error {
+			return sink.Close(ctx)
+		})
+	}
+	return g.Wait()
+}