clean up

[bytom/vapor.git] / toolbar / precog / monitor / monitor.go

diff --git a/toolbar/precog/monitor/monitor.go b/toolbar/precog/monitor/monitor.go
index 2d84587..b0de9a2 100644 (file)
--- a/toolbar/precog/monitor/monitor.go
+++ b/toolbar/precog/monitor/monitor.go
@@ -2,12 +2,10 @@ package monitor
 
 import (
        "fmt"
-       "io/ioutil"
        "os"
        "os/user"
        "strings"
        "sync"
-       "time"
 
        "github.com/jinzhu/gorm"
        log "github.com/sirupsen/logrus"
@@ -18,9 +16,8 @@ import (
        "github.com/vapor/event"
        "github.com/vapor/netsync/chainmgr"
        "github.com/vapor/netsync/consensusmgr"
-       "github.com/vapor/p2p"
-       // msgs "github.com/vapor/netsync/messages"
        "github.com/vapor/netsync/peers"
+       "github.com/vapor/p2p"
        "github.com/vapor/p2p/discover/dht"
        "github.com/vapor/p2p/discover/mdns"
        "github.com/vapor/p2p/signlib"
@@ -30,21 +27,23 @@ import (
 
 type monitor struct {
        *sync.RWMutex
-       cfg       *config.Config
-       db        *gorm.DB
-       nodeCfg   *vaporCfg.Config
-       sw        *p2p.Switch
-       discvCh   chan *dht.Node
-       privKey   chainkd.XPrv
-       chain     *mock.Chain
-       txPool    *mock.Mempool
-       connected uint32
+       cfg     *config.Config
+       db      *gorm.DB
+       nodeCfg *vaporCfg.Config
+       sw      *p2p.Switch
+       privKey chainkd.XPrv
+       chain   *mock.Chain
+       txPool  *mock.Mempool
+       // discvMap maps a node's public key to the node itself
+       discvMap      map[string]*dht.Node
+       dialCh        chan struct{}
+       checkStatusCh chan struct{}
 }
 
 // TODO: set myself as SPV?
 func NewMonitor(cfg *config.Config, db *gorm.DB) *monitor {
        //TODO: for test
-       cfg.CheckFreqSeconds = 120
+       cfg.CheckFreqSeconds = 15
 
        dbPath, err := makePath()
        if err != nil {
@@ -58,7 +57,6 @@ func NewMonitor(cfg *config.Config, db *gorm.DB) *monitor {
        }
        nodeCfg.DBPath = dbPath
        nodeCfg.ChainID = "mainnet"
-       discvCh := make(chan *dht.Node)
        privKey, err := signlib.NewPrivKey()
        if err != nil {
                log.Fatal(err)
@@ -70,14 +68,16 @@ func NewMonitor(cfg *config.Config, db *gorm.DB) *monitor {
        }
 
        return &monitor{
-               RWMutex: &sync.RWMutex{},
-               cfg:     cfg,
-               db:      db,
-               nodeCfg: nodeCfg,
-               discvCh: discvCh,
-               privKey: privKey.(chainkd.XPrv),
-               chain:   chain,
-               txPool:  txPool,
+               RWMutex:       &sync.RWMutex{},
+               cfg:           cfg,
+               db:            db,
+               nodeCfg:       nodeCfg,
+               privKey:       privKey.(chainkd.XPrv),
+               chain:         chain,
+               txPool:        txPool,
+               discvMap:      make(map[string]*dht.Node),
+               dialCh:        make(chan struct{}, 1),
+               checkStatusCh: make(chan struct{}, 1),
        }
 }
 
@@ -92,17 +92,10 @@ func makePath() (string, error) {
                return "", err
        }
 
-       tmpDir, err := ioutil.TempDir(dataPath, "")
-       if err != nil {
-               return "", err
-       }
-
-       return tmpDir, nil
+       return dataPath, nil
 }
 
 func (m *monitor) Run() {
-       defer os.RemoveAll(m.nodeCfg.DBPath)
-
        var seeds []string
        for _, node := range m.cfg.Nodes {
                seeds = append(seeds, fmt.Sprintf("%s:%d", node.Host, node.Port))
@@ -115,8 +108,8 @@ func (m *monitor) Run() {
                log.Fatal(err)
        }
 
+       m.dialCh <- struct{}{}
        go m.discoveryRoutine()
-       go m.collectDiscoveredNodes()
        go m.connectNodesRoutine()
        go m.checkStatusRoutine()
 }
@@ -151,7 +144,7 @@ func (m *monitor) prepareReactors(peers *peers.PeerSet) error {
        }
 
        for label, reactor := range m.sw.GetReactors() {
-               log.Debug("start reactor: (%s:%v)", label, reactor)
+               log.Debugf("start reactor: (%s:%v)", label, reactor)
                if _, err := reactor.Start(); err != nil {
                        return nil
                }
@@ -159,7 +152,6 @@ func (m *monitor) prepareReactors(peers *peers.PeerSet) error {
 
        m.sw.GetSecurity().RegisterFilter(m.sw.GetNodeInfo())
        m.sw.GetSecurity().RegisterFilter(m.sw.GetPeers())
-
        return m.sw.GetSecurity().Start()
 }
 
@@ -169,21 +161,13 @@ func (m *monitor) checkStatusRoutine() {
                log.Fatal(err)
        }
 
-       protocolReactor, ok := m.sw.GetReactors()["PROTOCOL"]
-       if !ok {
-               log.Fatal("protocolReactor not found")
-       }
-
        bestHeight := uint64(0)
-       ticker := time.NewTicker(time.Duration(m.cfg.CheckFreqSeconds) * time.Second)
-       for range ticker.C {
-               m.Lock()
-               log.Info("connected peers: ", m.sw.GetPeers().List())
-
+       for range m.checkStatusCh {
                for _, peer := range m.sw.GetPeers().List() {
                        peer.Start()
-                       protocolReactor.AddPeer(peer)
+                       peers.AddPeer(peer)
                }
+               log.Infof("%d connected peers: %v", len(m.sw.GetPeers().List()), m.sw.GetPeers().List())
 
                for _, peer := range m.sw.GetPeers().List() {
                        p := peers.GetPeer(peer.ID())
@@ -192,6 +176,7 @@ func (m *monitor) checkStatusRoutine() {
                        }
 
                        if err := p.SendStatus(m.chain.BestBlockHeader(), m.chain.LastIrreversibleHeader()); err != nil {
+                               log.Error(err)
                                peers.RemovePeer(p.ID())
                        }
                }
@@ -200,16 +185,9 @@ func (m *monitor) checkStatusRoutine() {
                        if peerInfo.Height > bestHeight {
                                bestHeight = peerInfo.Height
                        }
-
-                       m.savePeerInfo(peerInfo)
                }
                log.Info("bestHeight: ", bestHeight)
-
-               // TODO:
-               // msg := struct{ msgs.BlockchainMessage }{&msgs.GetBlockMessage{Height: bestHeight + 1}}
-               // for _, peer := range m.sw.GetPeers().List() {
-               //      peers.SendMsg(peer.ID(), msgs.BlockchainChannel, msg)
-               // }
+               m.processPeerInfos(peers.GetPeerInfos())
 
                for _, peer := range m.sw.GetPeers().List() {
                        p := peers.GetPeer(peer.ID())
@@ -221,5 +199,6 @@ func (m *monitor) checkStatusRoutine() {
                }
                log.Info("Disonnect all peers.")
                m.Unlock()
+               m.dialCh <- struct{}{}
        }
 }