Atomic syncs on connection states

This commit is contained in:
obscuren 2014-01-12 22:14:19 +01:00
parent f78bd4d5d0
commit 39bb2c94c0
2 changed files with 96 additions and 35 deletions

95
peer.go

@ -5,6 +5,8 @@ import (
"github.com/ethereum/ethwire-go" "github.com/ethereum/ethwire-go"
"log" "log"
"net" "net"
"sync/atomic"
"time"
) )
type Peer struct { type Peer struct {
@ -16,8 +18,12 @@ type Peer struct {
outputQueue chan *ethwire.InOutMsg outputQueue chan *ethwire.InOutMsg
// Quit channel // Quit channel
quit chan bool quit chan bool
// Determines whether it's an inbound or outbound peer
inbound bool // Determines whether it's an inbound or outbound peer inbound bool
// Flag for checking the peer's connectivity state
connected int32
disconnect int32
lastSend time.Time
} }
func NewPeer(conn net.Conn, server *Server, inbound bool) *Peer { func NewPeer(conn net.Conn, server *Server, inbound bool) *Peer {
@ -27,12 +33,57 @@ func NewPeer(conn net.Conn, server *Server, inbound bool) *Peer {
server: server, server: server,
conn: conn, conn: conn,
inbound: inbound, inbound: inbound,
disconnect: 0,
connected: 1,
} }
} }
func NewOutboundPeer(addr string, server *Server) *Peer {
p := &Peer{
outputQueue: make(chan *ethwire.InOutMsg, 1), // Buffered chan of 1 is enough
quit: make(chan bool),
server: server,
inbound: false,
connected: 0,
disconnect: 1,
}
// Set up the connection in another goroutine so we don't block the main thread
go func() {
conn, err := net.Dial("tcp", addr)
if err != nil {
p.Stop()
}
p.conn = conn
// Atomically set the connection state
atomic.StoreInt32(&p.connected, 1)
atomic.StoreInt32(&p.disconnect, 0)
log.Println("Connected to peer ::", conn.RemoteAddr())
}()
return p
}
// Outputs any RLP encoded data to the peer // Outputs any RLP encoded data to the peer
func (p *Peer) QueueMessage(msg *ethwire.InOutMsg) { func (p *Peer) QueueMessage(msg *ethwire.InOutMsg) {
p.outputQueue <- msg //ethwire.InOutMsg{MsgType: msgType, Nonce: ethutil.RandomUint64(), Data: data} p.outputQueue <- msg
}
func (p *Peer) writeMessage(msg *ethwire.InOutMsg) {
// Ignore the write if we're not connected
if atomic.LoadInt32(&p.connected) != 1 {
return
}
err := ethwire.WriteMessage(p.conn, msg)
if err != nil {
log.Println("Can't send message:", err)
// Stop the client if there was an error writing to it
p.Stop()
return
}
} }
// Outbound message handler. Outbound messages are handled here // Outbound message handler. Outbound messages are handled here
@ -42,28 +93,32 @@ out:
select { select {
// Main message queue. All outbound messages are processed through here // Main message queue. All outbound messages are processed through here
case msg := <-p.outputQueue: case msg := <-p.outputQueue:
// TODO Message checking and handle accordingly p.writeMessage(msg)
err := ethwire.WriteMessage(p.conn, msg)
if err != nil {
log.Println(err)
// Stop the client if there was an error writing to it
p.Stop()
}
p.lastSend = time.Now()
// Break out of the for loop if a quit message is posted // Break out of the for loop if a quit message is posted
case <-p.quit: case <-p.quit:
break out break out
} }
} }
clean:
// This loop is for draining the output queue and anybody waiting for us
for {
select {
case <- p.outputQueue:
// TODO
default:
break clean
}
}
} }
// Inbound handler. Inbound messages are received here and passed to the appropriate methods // Inbound handler. Inbound messages are received here and passed to the appropriate methods
func (p *Peer) HandleInbound() { func (p *Peer) HandleInbound() {
defer p.Stop()
out: out:
for { for atomic.LoadInt32(&p.disconnect) == 0 {
// Wait for a message from the peer // Wait for a message from the peer
msg, err := ethwire.ReadMessage(p.conn) msg, err := ethwire.ReadMessage(p.conn)
if err != nil { if err != nil {
@ -90,8 +145,7 @@ out:
} }
} }
// Notify the out handler we're quiting p.Stop()
p.quit <- true
} }
func (p *Peer) Start() { func (p *Peer) Start() {
@ -111,9 +165,16 @@ func (p *Peer) Start() {
} }
func (p *Peer) Stop() { func (p *Peer) Stop() {
p.conn.Close() if atomic.AddInt32(&p.disconnect, 1) != 1 {
return
}
p.quit <- true close(p.quit)
if atomic.LoadInt32(&p.connected) != 0 {
p.conn.Close()
}
log.Println("Peer shutdown")
} }
func (p *Peer) pushVersionAck() error { func (p *Peer) pushVersionAck() error {

@ -10,6 +10,16 @@ import (
"time" "time"
) )
func eachPeer(peers *list.List, callback func(*Peer)) {
// Loop thru the peers and close them (if we had them)
for e := peers.Front(); e != nil; e = e.Next() {
if peer, ok := e.Value.(*Peer); ok {
callback(peer)
}
}
}
type Server struct { type Server struct {
// Channel for shutting down the server // Channel for shutting down the server
shutdownChan chan bool shutdownChan chan bool
@ -57,27 +67,20 @@ func (s *Server) AddPeer(conn net.Conn) {
} }
func (s *Server) ConnectToPeer(addr string) error { func (s *Server) ConnectToPeer(addr string) error {
conn, err := net.Dial("tcp", addr) peer := NewOutboundPeer(addr, s)
if err != nil {
return err
}
peer := NewPeer(conn, s, false)
s.peers.PushBack(peer) s.peers.PushBack(peer)
peer.Start() peer.Start()
log.Println("Connected to peer ::", conn.RemoteAddr())
return nil return nil
} }
func (s *Server) Broadcast(msgType string, data []byte) { func (s *Server) Broadcast(msgType string, data []byte) {
for e := s.peers.Front(); e != nil; e = e.Next() { eachPeer(s.peers, func(p *Peer) {
if peer, ok := e.Value.(*Peer); ok { p.QueueMessage(ethwire.NewMessage(msgType, 0, data))
peer.QueueMessage(ethwire.NewMessage(msgType, 0, data)) })
}
}
} }
// Start the server // Start the server
@ -115,12 +118,9 @@ func (s *Server) Stop() {
// Close the database // Close the database
defer s.db.Close() defer s.db.Close()
// Loop thru the peers and close them (if we had them) eachPeer(s.peers, func(p *Peer) {
for e := s.peers.Front(); e != nil; e = e.Next() { p.Stop()
if peer, ok := e.Value.(*Peer); ok { })
peer.Stop()
}
}
s.shutdownChan <- true s.shutdownChan <- true
} }