pls, tx: Using in only transactionManager

tokamak-network · Apr 20, 2020 · 656e43d · 656e43d
1 parent a1dab59
commit 656e43d
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 12 deletions.
diff --git a/pls/rootchain_manager.go b/pls/rootchain_manager.go
@@ -69,7 +69,6 @@ type RootChainManager struct {
 	invalidExits map[uint64]map[uint64]invalidExits
 
 	// channels
-	wg 			 *sync.WaitGroup
 	quit             chan struct{}
 	epochPreparedCh  chan *rootchain.RootChainEpochPrepared
 	blockFinalizedCh chan *rootchain.RootChainBlockFinalized
@@ -130,7 +129,6 @@ func NewRootChainManager(
 		miner:             miner,
 		minerEnv:          env,
 		invalidExits:      make(map[uint64]map[uint64]invalidExits),
-		wg:				   new(sync.WaitGroup),
 		quit:              make(chan struct{}),
 		epochPreparedCh:   make(chan *rootchain.RootChainEpochPrepared, MAX_EPOCH_EVENTS),
 		blockFinalizedCh:  make(chan *rootchain.RootChainBlockFinalized),
@@ -156,7 +154,7 @@ func (rcm *RootChainManager) Start() error {
 	}
 
 	go rcm.pingBackend()
-	rcm.txManager.Start(rcm.wg)
+	rcm.txManager.Start()
 
 	if rcm.config.NodeMode == ModeOperator {
 		go rcm.miner.Start(rcm.config.Operator.Address, new(rootchain.RootChainEpochPrepared), true)
@@ -168,7 +166,6 @@ func (rcm *RootChainManager) Start() error {
 func (rcm *RootChainManager) Stop() error {
 	rcm.backend.Close()
 	rcm.txManager.Stop()
-	rcm.wg.Wait()
 	close(rcm.quit)
 	return nil
 }

diff --git a/tx/manager.go b/tx/manager.go
@@ -67,6 +67,7 @@ type TransactionManager struct {
 
 	lock         sync.RWMutex
 	gasPriceLock sync.Mutex
+	wg           *sync.WaitGroup
 	quit         chan struct{}
 }
 
@@ -256,8 +257,9 @@ func (tm *TransactionManager) Count(account accounts.Account, tx *types.Transact
 	return count
 }
 
-func (tm *TransactionManager) Start(wg *sync.WaitGroup) {
-	go tm.confirmLoop(wg)
+func (tm *TransactionManager) Start() {
+	tm.wg = new(sync.WaitGroup)
+	go tm.confirmLoop()
 
 	// send a single raw transaction to root chain.
 	// TODO: make it safe under root chain provider disconnect
@@ -421,9 +423,9 @@ func (tm *TransactionManager) Start(wg *sync.WaitGroup) {
 				}
 
 				for addr, _ := range tm.pending {
-					wg.Add(1)
+					tm.wg.Add(1)
 					go func(addr common.Address) {
-						defer wg.Done()
+						defer tm.wg.Done()
 						log.Trace("TransactionManager iterates", "addr", addr)
 						queue := tm.pending[addr]
 
@@ -707,10 +709,9 @@ func (tm *TransactionManager) indexOf(addr common.Address) int {
 }
 
 // TODO: use SubscribeNewHead with disconnection handling
-func (tm *TransactionManager) confirmLoop(wg *sync.WaitGroup) {
-	wg.Add(1)
-	defer wg.Done()
-
+func (tm *TransactionManager) confirmLoop() {
+	tm.wg.Add(1)
+	tm.wg.Done()
 	closed := false
 
 	newHeaderCh := make(chan *types.Header)
@@ -803,6 +804,7 @@ func (tm *TransactionManager) confirmLoop(wg *sync.WaitGroup) {
 }
 
 func (tm *TransactionManager) Stop() {
+	tm.wg.Wait()
 	close(tm.quit)
 }