OSDN Git Service

edit for code review
[bytom/bytom.git] / p2p / connection / connection.go
1 package connection
2
3 import (
4         "bufio"
5         "fmt"
6         "math"
7         "net"
8         "runtime/debug"
9         "sync/atomic"
10         "time"
11
12         log "github.com/sirupsen/logrus"
13         wire "github.com/tendermint/go-wire"
14         cmn "github.com/tendermint/tmlibs/common"
15         flow "github.com/tendermint/tmlibs/flowrate"
16 )
17
18 const (
19         packetTypePing           = byte(0x01)
20         packetTypePong           = byte(0x02)
21         packetTypeMsg            = byte(0x03)
22         maxMsgPacketPayloadSize  = 1024
23         maxMsgPacketOverheadSize = 10 // It's actually lower but good enough
24         maxMsgPacketTotalSize    = maxMsgPacketPayloadSize + maxMsgPacketOverheadSize
25
26         numBatchMsgPackets = 10
27         minReadBufferSize  = 1024
28         minWriteBufferSize = 65536
29         updateState        = 2 * time.Second
30         pingTimeout        = 40 * time.Second
31         flushThrottle      = 100 * time.Millisecond
32
33         defaultSendQueueCapacity   = 1
34         defaultSendRate            = int64(512000) // 500KB/s
35         defaultRecvBufferCapacity  = 4096
36         defaultRecvMessageCapacity = 22020096      // 21MB
37         defaultRecvRate            = int64(512000) // 500KB/s
38         defaultSendTimeout         = 10 * time.Second
39 )
40
41 type receiveCbFunc func(chID byte, msgBytes []byte)
42 type errorCbFunc func(interface{})
43
44 // Messages in channels are chopped into smaller msgPackets for multiplexing.
45 type msgPacket struct {
46         ChannelID byte
47         EOF       byte // 1 means message ends here.
48         Bytes     []byte
49 }
50
51 func (p msgPacket) String() string {
52         return fmt.Sprintf("MsgPacket{%X:%X T:%X}", p.ChannelID, p.Bytes, p.EOF)
53 }
54
55 /*
56 MConnection handles message transmission on multiple abstract communication
57 `Channel`s.  Each channel has a globally unique byte id.
58 The byte id and the relative priorities of each `Channel` are configured upon
59 initialization of the connection.
60
61 There are two methods for sending messages:
62         func (m MConnection) Send(chID byte, msg interface{}) bool {}
63         func (m MConnection) TrySend(chID byte, msg interface{}) bool {}
64
65 `Send(chID, msg)` is a blocking call that waits until `msg` is successfully queued
66 for the channel with the given id byte `chID`, or until the request times out.
67 The message `msg` is serialized using the `tendermint/wire` submodule's
68 `WriteBinary()` reflection routine.
69
70 `TrySend(chID, msg)` is a nonblocking call that returns false if the channel's
71 queue is full.
72
73 Inbound message bytes are handled with an onReceive callback function.
74 */
75 type MConnection struct {
76         cmn.BaseService
77
78         conn        net.Conn
79         bufReader   *bufio.Reader
80         bufWriter   *bufio.Writer
81         sendMonitor *flow.Monitor
82         recvMonitor *flow.Monitor
83         send        chan struct{}
84         pong        chan struct{}
85         channels    []*channel
86         channelsIdx map[byte]*channel
87         onReceive   receiveCbFunc
88         onError     errorCbFunc
89         errored     uint32
90         config      *MConnConfig
91
92         quit         chan struct{}
93         flushTimer   *cmn.ThrottleTimer // flush writes as necessary but throttled.
94         pingTimer    *time.Ticker       // send pings periodically
95         chStatsTimer *time.Ticker       // update channel stats periodically
96 }
97
98 // MConnConfig is a MConnection configuration.
99 type MConnConfig struct {
100         SendRate int64 `mapstructure:"send_rate"`
101         RecvRate int64 `mapstructure:"recv_rate"`
102 }
103
104 // DefaultMConnConfig returns the default config.
105 func DefaultMConnConfig() *MConnConfig {
106         return &MConnConfig{
107                 SendRate: defaultSendRate,
108                 RecvRate: defaultRecvRate,
109         }
110 }
111
112 // NewMConnectionWithConfig wraps net.Conn and creates multiplex connection with a config
113 func NewMConnectionWithConfig(conn net.Conn, chDescs []*ChannelDescriptor, onReceive receiveCbFunc, onError errorCbFunc, config *MConnConfig) *MConnection {
114         mconn := &MConnection{
115                 conn:        conn,
116                 bufReader:   bufio.NewReaderSize(conn, minReadBufferSize),
117                 bufWriter:   bufio.NewWriterSize(conn, minWriteBufferSize),
118                 sendMonitor: flow.New(0, 0),
119                 recvMonitor: flow.New(0, 0),
120                 send:        make(chan struct{}, 1),
121                 pong:        make(chan struct{}, 1),
122                 channelsIdx: map[byte]*channel{},
123                 channels:    []*channel{},
124                 onReceive:   onReceive,
125                 onError:     onError,
126                 config:      config,
127
128                 pingTimer:    time.NewTicker(pingTimeout),
129                 chStatsTimer: time.NewTicker(updateState),
130         }
131
132         for _, desc := range chDescs {
133                 descCopy := *desc // copy the desc else unsafe access across connections
134                 channel := newChannel(mconn, &descCopy)
135                 mconn.channelsIdx[channel.id] = channel
136                 mconn.channels = append(mconn.channels, channel)
137         }
138         mconn.BaseService = *cmn.NewBaseService(nil, "MConnection", mconn)
139         return mconn
140 }
141
142 // OnStart implements BaseService
143 func (c *MConnection) OnStart() error {
144         c.BaseService.OnStart()
145         c.quit = make(chan struct{})
146         c.flushTimer = cmn.NewThrottleTimer("flush", flushThrottle)
147         go c.sendRoutine()
148         go c.recvRoutine()
149         return nil
150 }
151
152 // OnStop implements BaseService
153 func (c *MConnection) OnStop() {
154         c.BaseService.OnStop()
155         c.flushTimer.Stop()
156         if c.quit != nil {
157                 close(c.quit)
158         }
159         c.conn.Close()
160         // We can't close pong safely here because recvRoutine may write to it after we've
161         // stopped. Though it doesn't need to get closed at all, we close it @ recvRoutine.
162 }
163
164 // CanSend returns true if you can send more data onto the chID, false otherwise
165 func (c *MConnection) CanSend(chID byte) bool {
166         if !c.IsRunning() {
167                 return false
168         }
169
170         channel, ok := c.channelsIdx[chID]
171         if !ok {
172                 return false
173         }
174         return channel.canSend()
175 }
176
177 // Send will queues a message to be sent to channel(blocking).
178 func (c *MConnection) Send(chID byte, msg interface{}) bool {
179         if !c.IsRunning() {
180                 return false
181         }
182
183         channel, ok := c.channelsIdx[chID]
184         if !ok {
185                 log.WithField("chID", chID).Error("cannot send bytes due to unknown channel")
186                 return false
187         }
188
189         if !channel.sendBytes(wire.BinaryBytes(msg)) {
190                 log.WithFields(log.Fields{"chID": chID, "conn": c, "msg": msg}).Error("MConnection send failed")
191                 return false
192         }
193
194         select {
195         case c.send <- struct{}{}:
196         default:
197         }
198         return true
199 }
200
201 // TrySend queues a message to be sent to channel(Nonblocking).
202 func (c *MConnection) TrySend(chID byte, msg interface{}) bool {
203         if !c.IsRunning() {
204                 return false
205         }
206
207         channel, ok := c.channelsIdx[chID]
208         if !ok {
209                 log.WithField("chID", chID).Error("cannot send bytes due to unknown channel")
210                 return false
211         }
212
213         ok = channel.trySendBytes(wire.BinaryBytes(msg))
214         if ok {
215                 select {
216                 case c.send <- struct{}{}:
217                 default:
218                 }
219         }
220         return ok
221 }
222
223 func (c *MConnection) String() string {
224         return fmt.Sprintf("MConn{%v}", c.conn.RemoteAddr())
225 }
226
227 func (c *MConnection) flush() {
228         if err := c.bufWriter.Flush(); err != nil {
229                 log.WithField("error", err).Error("MConnection flush failed")
230         }
231 }
232
233 // Catch panics, usually caused by remote disconnects.
234 func (c *MConnection) _recover() {
235         if r := recover(); r != nil {
236                 stack := debug.Stack()
237                 err := cmn.StackError{r, stack}
238                 c.stopForError(err)
239         }
240 }
241
242 // recvRoutine reads msgPackets and reconstructs the message using the channels' "recving" buffer.
243 // After a whole message has been assembled, it's pushed to onReceive().
244 // Blocks depending on how the connection is throttled.
245 func (c *MConnection) recvRoutine() {
246         defer c._recover()
247         defer close(c.pong)
248
249         for {
250                 // Block until .recvMonitor says we can read.
251                 c.recvMonitor.Limit(maxMsgPacketTotalSize, atomic.LoadInt64(&c.config.RecvRate), true)
252
253                 // Read packet type
254                 var n int
255                 var err error
256                 pktType := wire.ReadByte(c.bufReader, &n, &err)
257                 c.recvMonitor.Update(int(n))
258                 if err != nil {
259                         if c.IsRunning() {
260                                 log.WithFields(log.Fields{"conn": c, "error": err}).Error("Connection failed @ recvRoutine (reading byte)")
261                                 c.conn.Close()
262                                 c.stopForError(err)
263                         }
264                         return
265                 }
266
267                 // Read more depending on packet type.
268                 switch pktType {
269                 case packetTypePing:
270                         log.Debug("receive Ping")
271                         select {
272                         case c.pong <- struct{}{}:
273                         default:
274                         }
275
276                 case packetTypePong:
277                         log.Debug("receive Pong")
278
279                 case packetTypeMsg:
280                         pkt, n, err := msgPacket{}, int(0), error(nil)
281                         wire.ReadBinaryPtr(&pkt, c.bufReader, maxMsgPacketTotalSize, &n, &err)
282                         c.recvMonitor.Update(int(n))
283                         if err != nil {
284                                 if c.IsRunning() {
285                                         log.WithFields(log.Fields{"conn": c, "error": err}).Error("failed on recvRoutine")
286                                         c.stopForError(err)
287                                 }
288                                 return
289                         }
290
291                         channel, ok := c.channelsIdx[pkt.ChannelID]
292                         if !ok || channel == nil {
293                                 cmn.PanicQ(cmn.Fmt("Unknown channel %X", pkt.ChannelID))
294                         }
295
296                         msgBytes, err := channel.recvMsgPacket(pkt)
297                         if err != nil {
298                                 if c.IsRunning() {
299                                         log.WithFields(log.Fields{"conn": c, "error": err}).Error("failed on recvRoutine")
300                                         c.stopForError(err)
301                                 }
302                                 return
303                         }
304
305                         if msgBytes != nil {
306                                 c.onReceive(pkt.ChannelID, msgBytes)
307                         }
308
309                 default:
310                         cmn.PanicSanity(cmn.Fmt("Unknown message type %X", pktType))
311                 }
312         }
313 }
314
315 // Returns true if messages from channels were exhausted.
316 func (c *MConnection) sendMsgPacket() bool {
317         var leastRatio float32 = math.MaxFloat32
318         var leastChannel *channel
319         for _, channel := range c.channels {
320                 if !channel.isSendPending() {
321                         continue
322                 }
323                 if ratio := float32(channel.recentlySent) / float32(channel.priority); ratio < leastRatio {
324                         leastRatio = ratio
325                         leastChannel = channel
326                 }
327         }
328         if leastChannel == nil {
329                 return true
330         }
331
332         n, err := leastChannel.writeMsgPacketTo(c.bufWriter)
333         if err != nil {
334                 log.WithField("error", err).Error("failed to write msgPacket")
335                 c.stopForError(err)
336                 return true
337         }
338         c.sendMonitor.Update(int(n))
339         c.flushTimer.Set()
340         return false
341 }
342
343 // sendRoutine polls for packets to send from channels.
344 func (c *MConnection) sendRoutine() {
345         defer c._recover()
346
347         for {
348                 var n int
349                 var err error
350                 select {
351                 case <-c.flushTimer.Ch:
352                         c.flush()
353                 case <-c.chStatsTimer.C:
354                         for _, channel := range c.channels {
355                                 channel.updateStats()
356                         }
357                 case <-c.pingTimer.C:
358                         log.Debug("send Ping")
359                         wire.WriteByte(packetTypePing, c.bufWriter, &n, &err)
360                         c.sendMonitor.Update(int(n))
361                         c.flush()
362                 case <-c.pong:
363                         log.Debug("send Pong")
364                         wire.WriteByte(packetTypePong, c.bufWriter, &n, &err)
365                         c.sendMonitor.Update(int(n))
366                         c.flush()
367                 case <-c.quit:
368                         return
369                 case <-c.send:
370                         if eof := c.sendSomeMsgPackets(); !eof {
371                                 select {
372                                 case c.send <- struct{}{}:
373                                 default:
374                                 }
375                         }
376                 }
377
378                 if !c.IsRunning() {
379                         return
380                 }
381                 if err != nil {
382                         log.WithFields(log.Fields{"conn": c, "error": err}).Error("Connection failed @ sendRoutine")
383                         c.stopForError(err)
384                         return
385                 }
386         }
387 }
388
389 // Returns true if messages from channels were exhausted.
390 func (c *MConnection) sendSomeMsgPackets() bool {
391         // Block until .sendMonitor says we can write.
392         // Once we're ready we send more than we asked for,
393         // but amortized it should even out.
394         c.sendMonitor.Limit(maxMsgPacketTotalSize, atomic.LoadInt64(&c.config.SendRate), true)
395         for i := 0; i < numBatchMsgPackets; i++ {
396                 if c.sendMsgPacket() {
397                         return true
398                 }
399         }
400         return false
401 }
402
403 func (c *MConnection) stopForError(r interface{}) {
404         c.Stop()
405         if atomic.CompareAndSwapUint32(&c.errored, 0, 1) && c.onError != nil {
406                 c.onError(r)
407         }
408 }