Merge a89b28f9b4 into dc7f202ecd

common/shutdown.go

@@ -0,0 +1,79 @@
+// Copyright 2014 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
+
+package common
+
+import (
+	"sync"
+)
+
+// ShutdownManager implements a sync mechanism that waits for multiple jobs
+// to finish and also prevents starting new jobs after the shutdown process
+// has started.
+type ShutdownManager struct {
+	wg sync.WaitGroup
+	lock sync.RWMutex
+	stop bool
+	stopChn chan struct{}
+}
+
+// NewShutdownManager creates a new ShutdownManager instance
+func NewShutdownManager() *ShutdownManager {
+	return &ShutdownManager{stopChn: make(chan struct{})}
+}
+
+// Enter should be called before beginning a new job. If Enter returns false,
+// the shutdown process has already been started. In this case, Exit should not
+// be called.
+func (self *ShutdownManager) Enter() bool {
+	self.lock.RLock()
+	defer self.lock.RUnlock()
+
+	if self.stop {
+		return false
+	}
+	self.wg.Add(1)
+	return true
+}
+
+// Exit should be called after finishing a job
+func (self *ShutdownManager) Exit() {
+	self.wg.Done()
+}
+
+// Shutdown initiates the shutdown process. After calling it, any subsequent
+// calls to Enter will return false. Shutdown returns after all active jobs
+// have finished and called Exit.
+func (self *ShutdownManager) Shutdown() {
+	self.lock.Lock()
+	self.stop = true
+	self.lock.Unlock()
+	close(self.stopChn)
+	self.wg.Wait()
+}
+
+// Stopped returns true if the shutdown process has already been started
+func (self *ShutdownManager) Stopped() bool {
+	self.lock.RLock()
+	defer self.lock.RUnlock()
+	
+	return self.stop
+}
+
+// StopChannel returns a channel which is closed when the shutdown starts
+func (self *ShutdownManager) StopChannel() chan struct{} {
+	return self.stopChn
+}

core/tx_pool.go

@@ -99,6 +99,10 @@ func (pool *TxPool) eventLoop() {
 	// Track chain events. When a chain events occurs (new chain canon block)
 	// we need to know the new state. The new state will help us determine
 	// the nonces in the managed state
+	if !pool.events.LoopStarted() {
+		return
+	}
+	defer pool.events.LoopStopped()
 	for ev := range pool.events.Chan() {
 		switch ev := ev.Data.(type) {
 		case ChainHeadEvent:

eth/filters/filter_system.go

@@ -131,6 +131,10 @@ func (fs *FilterSystem) Get(id int) *Filter {
 // filterLoop waits for specific events from ethereum and fires their handlers
 // when the filter matches the requirements.
 func (fs *FilterSystem) filterLoop() {
+	if !fs.sub.LoopStarted() {
+		return
+	}
+	defer fs.sub.LoopStopped()
 	for event := range fs.sub.Chan() {
 		switch ev := event.Data.(type) {
 		case core.ChainEvent:

eth/gasprice.go

@@ -103,6 +103,10 @@ func (self *GasPriceOracle) processPastBlocks(chain *core.BlockChain) {
 func (self *GasPriceOracle) listenLoop() {
 	events := self.eth.EventMux().Subscribe(core.ChainEvent{}, core.ChainSplitEvent{})
 	defer events.Unsubscribe()
+	if !events.LoopStarted() {
+		return
+	}
+	defer events.LoopStopped()
 
 	for event := range events.Chan() {
 		switch event := event.Data.(type) {

eth/handler.go

@@ -82,6 +82,7 @@ type ProtocolManager struct {
 	// and processing
 	wg   sync.WaitGroup
 	quit bool
+	sm   *common.ShutdownManager
 }
 
 // NewProtocolManager returns a new ethereum sub protocol manager. The Ethereum sub protocol manages peers capable
@@ -104,6 +105,7 @@ func NewProtocolManager(config *core.ChainConfig, fastSync bool, networkId int,
 		newPeerCh:  make(chan *peer, 1),
 		txsyncCh:   make(chan *txsync),
 		quitSync:   make(chan struct{}),
+		sm:			common.NewShutdownManager(),
 	}
 	// Initiate a sub-protocol for every implemented version we can handle
 	manager.SubProtocols = make([]p2p.Protocol, 0, len(ProtocolVersions))
@@ -119,9 +121,17 @@ func NewProtocolManager(config *core.ChainConfig, fastSync bool, networkId int,
 			Version: version,
 			Length:  ProtocolLengths[i],
 			Run: func(p *p2p.Peer, rw p2p.MsgReadWriter) error {
+				if !manager.sm.Enter() {
+					return nil
+				}
+				defer manager.sm.Exit()
 				peer := manager.newPeer(int(version), p, rw)
-				manager.newPeerCh <- peer
+				select {
-				return manager.handle(peer)
+				case manager.newPeerCh <- peer:
+					return manager.handle(peer)
+				case <-manager.sm.StopChannel():
+					return nil
+				}
 			},
 			NodeInfo: func() interface{} {
 				return manager.NodeInfo()
@@ -196,8 +206,25 @@ func (pm *ProtocolManager) Stop() {
 	pm.minedBlockSub.Unsubscribe() // quits blockBroadcastLoop
 	close(pm.quitSync)             // quits syncer, fetcher, txsyncLoop
 
+	stopped := make(chan struct{})
+	go func() {
+		for {
+			// drop all peers
+			for _, id := range pm.peers.AllPeerIDs() {
+				pm.removePeer(id)
+			}
+			select {
+				case <-stopped:
+					return
+				case <-time.After(time.Millisecond * 100):
+			}
+		}
+	}()
+
 	// Wait for any process action
+	pm.sm.Shutdown()
 	pm.wg.Wait()
+	close(stopped)
 
 	glog.V(logger.Info).Infoln("Ethereum protocol handler stopped")
 }
@@ -239,7 +266,7 @@ func (pm *ProtocolManager) handle(p *peer) error {
 	pm.syncTransactions(p)
 
 	// main loop. handle incoming messages.
-	for {
+	for !pm.sm.Stopped() {
 		if err := pm.handleMsg(p); err != nil {
 			glog.V(logger.Debug).Infof("%v: message handling failed: %v", p, err)
 			return err
@@ -255,6 +282,9 @@ func (pm *ProtocolManager) handleMsg(p *peer) error {
 	if err != nil {
 		return err
 	}
+	if pm.sm.Stopped() {
+		return nil
+	}
 	if msg.Size > ProtocolMaxMsgSize {
 		return errResp(ErrMsgTooLarge, "%v > %v", msg.Size, ProtocolMaxMsgSize)
 	}
@@ -733,6 +763,10 @@ func (pm *ProtocolManager) BroadcastTx(hash common.Hash, tx *types.Transaction)
 // Mined broadcast loop
 func (self *ProtocolManager) minedBroadcastLoop() {
 	// automatically stops if unsubscribe
+	if !self.minedBlockSub.LoopStarted() {
+		return
+	}
+	defer self.minedBlockSub.LoopStopped()
 	for obj := range self.minedBlockSub.Chan() {
 		switch ev := obj.Data.(type) {
 		case core.NewMinedBlockEvent:
@@ -744,6 +778,10 @@ func (self *ProtocolManager) minedBroadcastLoop() {
 
 func (self *ProtocolManager) txBroadcastLoop() {
 	// automatically stops if unsubscribe
+	if !self.txSub.LoopStarted() {
+		return
+	}
+	defer self.txSub.LoopStopped()
 	for obj := range self.txSub.Chan() {
 		event := obj.Data.(core.TxPreEvent)
 		self.BroadcastTx(event.Tx.Hash(), event.Tx)

eth/peer.go

@@ -388,6 +388,20 @@ func (ps *peerSet) Unregister(id string) error {
 	return nil
 }
 
+// AllPeerIDs returns a list of all registered peer IDs
+func (ps *peerSet) AllPeerIDs() []string {
+	ps.lock.RLock()
+	defer ps.lock.RUnlock()
+	
+	res := make([]string, len(ps.peers))	
+	idx := 0
+	for id, _ := range ps.peers {
+		res[idx] = id
+		idx++
+	}
+	return res
+}
+
 // Peer retrieves the registered peer with the given id.
 func (ps *peerSet) Peer(id string) *peer {
 	ps.lock.RLock()

event/event.go

@@ -23,6 +23,8 @@ import (
 	"reflect"
 	"sync"
 	"time"
+
+	"github.com/ethereum/go-ethereum/common"
 )
 
 // Event is a time-tagged notification pushed to subscribers.
@@ -42,6 +44,8 @@ type Subscription interface {
 	// The event channel is closed.
 	// Unsubscribe can be called more than once.
 	Unsubscribe()
+	LoopStarted() bool
+	LoopStopped()
 }
 
 // A TypeMux dispatches events to registered receivers. Receivers can be
@@ -53,6 +57,7 @@ type TypeMux struct {
 	mutex   sync.RWMutex
 	subm    map[reflect.Type][]*muxsub
 	stopped bool
+	sm 		*common.ShutdownManager
 }
 
 // ErrMuxClosed is returned when Posting on a closed TypeMux.
@@ -68,6 +73,9 @@ func (mux *TypeMux) Subscribe(types ...interface{}) Subscription {
 	if mux.stopped {
 		close(sub.postC)
 	} else {
+		if mux.sm == nil {
+			mux.sm = common.NewShutdownManager()
+		}
 		if mux.subm == nil {
 			mux.subm = make(map[reflect.Type][]*muxsub)
 		}
@@ -119,7 +127,11 @@ func (mux *TypeMux) Stop() {
 	}
 	mux.subm = nil
 	mux.stopped = true
+	sm := mux.sm
 	mux.mutex.Unlock()
+	if sm != nil {
+		sm.Shutdown()
+	}
 }
 
 func (mux *TypeMux) del(s *muxsub) {
@@ -187,6 +199,14 @@ func (s *muxsub) Unsubscribe() {
 	s.closewait()
 }
 
+func (s *muxsub) LoopStarted() bool {
+	return s.mux.sm.Enter()
+}
+
+func (s *muxsub) LoopStopped() {
+	s.mux.sm.Exit()
+}
+
 func (s *muxsub) closewait() {
 	s.closeMu.Lock()
 	defer s.closeMu.Unlock()
@@ -215,4 +235,4 @@ func (s *muxsub) deliver(event *Event) {
 	case s.postC <- event:
 	case <-s.closing:
 	}
-}
+}

miner/miner.go

@@ -64,6 +64,10 @@ func New(eth core.Backend, config *core.ChainConfig, mux *event.TypeMux, pow pow
 // and halt your mining operation for as long as the DOS continues.
 func (self *Miner) update() {
 	events := self.mux.Subscribe(downloader.StartEvent{}, downloader.DoneEvent{}, downloader.FailedEvent{})
+	if !events.LoopStarted() {
+		return
+	}
+	defer events.LoopStopped()
 out:
 	for ev := range events.Chan() {
 		switch ev.Data.(type) {

miner/worker.go

@@ -220,6 +220,9 @@ func (self *worker) update() {
 	eventSub := self.mux.Subscribe(core.ChainHeadEvent{}, core.ChainSideEvent{}, core.TxPreEvent{})
 	defer eventSub.Unsubscribe()
 
+	if !eventSub.LoopStarted() {
+		return
+	}
 	eventCh := eventSub.Chan()
 	for {
 		select {
@@ -227,6 +230,7 @@ func (self *worker) update() {
 			if !ok {
 				// Event subscription closed, set the channel to nil to stop spinning
 				eventCh = nil
+				eventSub.LoopStopped()
 				continue
 			}
 			// A real event arrived, process interesting content