Prepare for multishard tx

Author: reshke

router/frontend/frontend.go

@@ -38,6 +38,7 @@ func ProcessMessage(qr qrouter.QueryRouter, rst relay.RelayStateMgr, msg pgproto
 				rst.AddQuery(q)
 				return rst.ProcessMessageBuf(true, true)
 			}, true); err != nil {
+				_ = rst.CompleteRelay(true)
 				return err
 			}
 
@@ -66,6 +67,7 @@ func ProcessMessage(qr qrouter.QueryRouter, rst relay.RelayStateMgr, msg pgproto
 
 				return rst.ProcessMessageBuf(true, true)
 			}, false); err != nil {
+				_ = rst.CompleteRelay(true)
 				return err
 			}
 
@@ -132,6 +134,7 @@ func ProcessMessage(qr qrouter.QueryRouter, rst relay.RelayStateMgr, msg pgproto
 			// this call compeletes relay, sends RFQ
 			return rst.ProcessMessageBuf(true, true)
 		}, false); err != nil {
+			_ = rst.CompleteRelay(true)
 			return err
 		}
 

router/frontend/frontend_test.go

@@ -179,6 +179,8 @@ func TestFrontendXProto(t *testing.T) {
 		PoolMode: config.PoolModeTransaction,
 	}
 
+	qr.EXPECT().Mgr().Return(mmgr).AnyTimes()
+
 	beRule := &config.BackendRule{}
 
 	qr.EXPECT().Mgr().Return(mmgr).AnyTimes()

router/relay/phandler.go

@@ -10,6 +10,7 @@ import (
 	"github.com/pg-sharding/spqr/router/client"
 	"github.com/pg-sharding/spqr/router/parser"
 	"github.com/pg-sharding/spqr/router/pgcopy"
+	"github.com/pg-sharding/spqr/router/server"
 )
 
 // Execute requered command via
@@ -19,6 +20,8 @@ type QueryStateExecutor interface {
 
 	Client() client.RouterClient
 
+	Deploy(server server.Server) error
+
 	ExecBegin(rst RelayStateMgr, query string, st *parser.ParseStateTXBegin) error
 	ExecCommit(rst RelayStateMgr, query string) error
 	ExecRollback(rst RelayStateMgr, query string) error

router/relay/phs.go

@@ -18,6 +18,7 @@ import (
 	"github.com/pg-sharding/spqr/router/parser"
 	"github.com/pg-sharding/spqr/router/pgcopy"
 	"github.com/pg-sharding/spqr/router/rmeta"
+	"github.com/pg-sharding/spqr/router/server"
 
 	"github.com/pg-sharding/lyx/lyx"
 )
@@ -27,6 +28,24 @@ type QueryStateExecutorImpl struct {
 
 	txStatus txstatus.TXStatus
 	cl       client.RouterClient
+
+	savedBegin *pgproto3.Query
+}
+
+// Deploy implements QueryStateExecutor.
+func (s *QueryStateExecutorImpl) Deploy(server server.Server) error {
+	if s.txStatus == txstatus.TXIDLE {
+		return nil
+	}
+
+	for _, sh := range server.Datashards() {
+		if err := sh.Send(s.savedBegin); err != nil {
+			s.txStatus = txstatus.TXERR
+			return err
+		}
+	}
+
+	return nil
 }
 
 func (s *QueryStateExecutorImpl) SetTxStatus(status txstatus.TXStatus) {
@@ -41,12 +60,10 @@ func (s *QueryStateExecutorImpl) TxStatus() txstatus.TXStatus {
 
 func (s *QueryStateExecutorImpl) ExecBegin(rst RelayStateMgr, query string, st *parser.ParseStateTXBegin) error {
 	// explicitly set silent query message, as it can differ from query begin in xporot
-	rst.AddSilentQuery(&pgproto3.Query{
-		String: query,
-	})
 
 	s.SetTxStatus(txstatus.TXACT)
 	s.cl.StartTx()
+	s.savedBegin = &pgproto3.Query{String: query}
 
 	spqrlog.Zero.Debug().Msg("start new transaction")
 
@@ -59,7 +76,6 @@ func (s *QueryStateExecutorImpl) ExecBegin(rst RelayStateMgr, query string, st *
 		}
 	}
 	return rst.Client().ReplyCommandComplete("BEGIN")
-
 }
 
 // query in commit query. maybe commit or commit `name`
@@ -73,14 +89,16 @@ func (s *QueryStateExecutorImpl) ExecCommit(rst RelayStateMgr, query string) err
 		rst.Flush()
 		return nil
 	}
-	rst.AddQuery(&pgproto3.Query{
-		String: query,
-	})
-	err := rst.ProcessMessageBuf(true, true)
-	if err == nil {
-		rst.Client().CommitActiveSet()
+
+	for _, sh := range rst.Client().Server().Datashards() {
+		if err := sh.Send(&pgproto3.Query{String: query}); err != nil {
+			s.txStatus = txstatus.TXERR
+			return err
+		}
 	}
-	return err
+
+	rst.Client().CommitActiveSet()
+	return nil
 }
 
 /* TODO: proper support for rollback to savepoint */
@@ -94,14 +112,16 @@ func (s *QueryStateExecutorImpl) ExecRollback(rst RelayStateMgr, query string) e
 		rst.Flush()
 		return nil
 	}
-	rst.AddQuery(&pgproto3.Query{
-		String: query,
-	})
-	err := rst.ProcessMessageBuf(true, true)
-	if err == nil {
-		s.cl.Rollback()
+
+	for _, sh := range rst.Client().Server().Datashards() {
+		if err := sh.Send(&pgproto3.Query{String: query}); err != nil {
+			s.txStatus = txstatus.TXERR
+			return err
+		}
 	}
-	return err
+
+	s.cl.Rollback()
+	return nil
 }
 
 func (s *QueryStateExecutorImpl) ExecSet(rst RelayStateMgr, query string, name, value string) error {

router/relay/relay.go

@@ -501,7 +501,7 @@ var ErrSkipQuery = fmt.Errorf("wait for a next query")
 
 // TODO : unit tests
 func (rst *RelayStateImpl) procRoutes(routes []*kr.ShardKey) error {
-	// if there is no routes configurted, there is nowhere to route to
+	// if there is no routes configured, there is nowhere to route to
 	if len(routes) == 0 {
 		return qrouter.MatchShardError
 	}
@@ -531,6 +531,15 @@ func (rst *RelayStateImpl) procRoutes(routes []*kr.ShardKey) error {
 			Uint("client", rst.Client().ID()).
 			Msg("client encounter while initialing server connection")
 
+		_ = rst.Unroute(rst.activeShards)
+		_ = rst.Reset()
+		return err
+	}
+
+	/* if transaction is explicitly requested, deploy */
+
+	if err := rst.QueryExecutor().Deploy(rst.Client().Server()); err != nil {
+		_ = rst.Unroute(rst.activeShards)
 		_ = rst.Reset()
 		return err
 	}
@@ -756,6 +765,10 @@ func (rst *RelayStateImpl) CompleteRelay(replyCl bool) error {
 			return nil
 		}
 
+		_ = rst.Reset()
+
+		_ = rst.Unroute(rst.activeShards)
+
 		if replyCl {
 			if err := rst.Cl.Send(&pgproto3.ReadyForQuery{
 				TxStatus: byte(txstatus.TXIDLE),
@@ -793,7 +806,10 @@ func (rst *RelayStateImpl) CompleteRelay(replyCl bool) error {
 				return err
 			}
 		}
-		return nil
+
+		_ = rst.Reset()
+
+		return rst.Unroute(rst.activeShards)
 	default:
 		err := fmt.Errorf("unknown tx status %v", rst.qse.TxStatus())
 		return err
@@ -1087,6 +1103,7 @@ func (rst *RelayStateImpl) ProcessExtendedBuffer() error {
 
 				return nil
 			}, true /* cache parsing for prep statement */); err != nil {
+				_ = rst.CompleteRelay(true)
 				return err
 			}