tpch: check connection before each run iteration (#147)

Signed-off-by: zyguan <[email protected]> Signed-off-by: zyguan <[email protected]>
pingcap · Oct 9, 2022 · b1e0fb5 · b1e0fb5
1 parent c482fcf
commit b1e0fb5
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/pkg/workload/base.go b/pkg/workload/base.go
@@ -20,6 +20,9 @@ type TpcState struct {
 }
 
 func (t *TpcState) RefreshConn(ctx context.Context) error {
+	if t.Conn != nil {
+		t.Conn.Close()
+	}
 	conn, err := t.DB.Conn(ctx)
 	if err != nil {
 		return err

diff --git a/tpch/workload.go b/tpch/workload.go
@@ -201,6 +201,11 @@ func (w *Workloader) CheckPrepare(ctx context.Context, threadID int) error {
 func (w *Workloader) Run(ctx context.Context, threadID int) error {
 	s := w.getState(ctx)
 	defer w.updateState(ctx)
+	if err := s.Conn.PingContext(ctx); err != nil {
+		if err := s.RefreshConn(ctx); err != nil {
+			return err
+		}
+	}
 
 	queryName := w.cfg.QueryNames[s.queryIdx%len(w.cfg.QueryNames)]
 	query := query(w.cfg.Driver, queryName)