tpcc: refresh conn if pingContext fails (#76)

Signed-off-by: mahjonp <junpeng.man@gmail.com>

pkg/workload/base.go

@@ -11,13 +11,23 @@ import (
 
 // TpcState saves state for each thread
 type TpcState struct {
+	DB   *sql.DB
 	Conn *sql.Conn
 
 	R *rand.Rand
 
 	Buf *util.BufAllocator
 }
 
+func (t *TpcState) RefreshConn(ctx context.Context) error {
+	conn, err := t.DB.Conn(ctx)
+	if err != nil {
+		return err
+	}
+	t.Conn = conn
+	return nil
+}
+
 // NewTpcState creates a base TpcState
 func NewTpcState(ctx context.Context, db *sql.DB) *TpcState {
 	var conn *sql.Conn
@@ -32,6 +42,7 @@ func NewTpcState(ctx context.Context, db *sql.DB) *TpcState {
 	r := rand.New(rand.NewSource(time.Now().UnixNano()))
 
 	s := &TpcState{
+		DB:   db,
 		Conn: conn,
 		R:    r,
 		Buf:  util.NewBufAllocator(),

tpcc/workload.go

@@ -192,8 +192,14 @@ func getTPCCState(ctx context.Context) *tpccState {
 // Run implements Workloader interface
 func (w *Workloader) Run(ctx context.Context, threadID int) error {
 	s := getTPCCState(ctx)
-
-	if s.newOrderStmts == nil {
+	refreshConn := false
+	if err := s.Conn.PingContext(ctx); err != nil {
+		if err := s.RefreshConn(ctx); err != nil {
+			return err
+		}
+		refreshConn = true
+	}
+	if s.newOrderStmts == nil || refreshConn {
 		s.newOrderStmts = map[string]*sql.Stmt{
 			newOrderSelectCustomer: prepareStmt(ctx, s.Conn, newOrderSelectCustomer),
 			newOrderSelectDistrict: prepareStmt(ctx, s.Conn, newOrderSelectDistrict),