├── .gitignore ├── LICENSE ├── README.md ├── client.go ├── doc.go ├── example_test.go ├── io.go ├── pumper.go └── server.go /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2015, someonegg 2 | All rights reserved. 3 | 4 | Redistribution and use in source and binary forms, with or without 5 | modification, are permitted provided that the following conditions are met: 6 | 7 | * Redistributions of source code must retain the above copyright notice, this 8 | list of conditions and the following disclaimer. 9 | 10 | * Redistributions in binary form must reproduce the above copyright notice, 11 | this list of conditions and the following disclaimer in the documentation 12 | and/or other materials provided with the distribution. 13 | 14 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 15 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 16 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 17 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 18 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 19 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 20 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 21 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 22 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 23 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 24 | 25 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | bdmsg 2 | ====== 3 | 4 | bdmsg implements bidirectional directly message protocol with golang. 5 | Message is defined as variable-length byte array, they are 6 | distinguished by message-type, they can be freely transferred between 7 | the server and the client. 8 | 9 | This package now has a better alternative which is more concise and customizable: 10 | [msgpump](https://github.com/someonegg/msgpump). 11 | 12 | Documentation 13 | ------------- 14 | 15 | - [API Reference](http://godoc.org/github.com/someonegg/bdmsg) 16 | - [Example](https://github.com/someonegg/bdmsg-example.git) 17 | 18 | Installation 19 | ------------ 20 | 21 | Install bdmsg using the "go get" command: 22 | 23 | go get github.com/someonegg/bdmsg 24 | 25 | The Go distribution is bdmsg's only dependency. 26 | -------------------------------------------------------------------------------- /client.go: -------------------------------------------------------------------------------- 1 | // Copyright 2015 someonegg. All rights reserved. 2 | // Use of this source code is governed by a BSD-style 3 | // license that can be found in the LICENSE file. 4 | 5 | package bdmsg 6 | 7 | import ( 8 | "golang.org/x/net/context" 9 | "net" 10 | ) 11 | 12 | // Client represents a message client in the client side. 13 | // You can send(recive) message to(from) server use Client.Pumper(). 14 | // 15 | // Multiple goroutines can invoke methods on a Client simultaneously. 16 | type Client struct { 17 | Pumper 18 | c net.Conn 19 | } 20 | 21 | // NewClient allocates and returns a new Client. 22 | // 23 | // The ownership of c will be transferred to Client, dont 24 | // control it in other places. 25 | func NewClient(parent context.Context, c net.Conn, ioc Converter, 26 | h PumperHandler, pumperInN, pumperOutN int) *Client { 27 | 28 | rw := ioc.Convert(c) 29 | 30 | t := &Client{} 31 | t.c = c 32 | t.Pumper.init(rw, h, pumperInN, pumperOutN) 33 | t.Pumper.SetUserData(t) 34 | t.Pumper.Start(parent, t) 35 | 36 | return t 37 | } 38 | 39 | func (c *Client) OnStop() { 40 | c.c.Close() 41 | if sn, ok := c.rw.(StopNotifier); ok { 42 | sn.OnStop() 43 | } 44 | } 45 | 46 | // InnerPumper is c.Pumper, its initial userdata is *Client. 47 | func (c *Client) InnerPumper() *Pumper { 48 | return &c.Pumper 49 | } 50 | 51 | func (c *Client) Conn() net.Conn { 52 | return c.c 53 | } 54 | -------------------------------------------------------------------------------- /doc.go: -------------------------------------------------------------------------------- 1 | // Copyright 2015 someonegg. All rights reserved. 2 | // Use of this source code is governed by a BSD-style 3 | // license that can be found in the LICENSE file. 4 | 5 | // Package bdmsg implements bidirectional directly message protocol. 6 | // 7 | // This package now has a better alternative which is more concise and customizable: 8 | // 9 | // https://github.com/someonegg/msgpump 10 | // 11 | // Message is defined as variable-length byte array, they are 12 | // distinguished by message-type, they can be freely transferred between 13 | // the server and the client. 14 | // 15 | // You can provide your own "Converter" to customize the message IO. 16 | // The default converter "DefaultIOC" use "MsgRWIO", message layout is: 17 | // A(length) + B(type) + C(data) 18 | // A is 4-bytes int, big-endian 19 | // B is 4-bytes int, big-endian 20 | // C is byte array, its length is (A) - 4 21 | // 22 | // Here is a quick example, includes protocol, server, client. 23 | // 24 | // Protocol 25 | // const ( 26 | // MsgTypeConnect = 1 27 | // MsgTypeConnectReply = 2 28 | // ) 29 | // 30 | // type ConnectRequst struct { 31 | // Name string 32 | // Pass string 33 | // } 34 | // 35 | // func (p *ConnectRequst) Marshal() ([]byte, error) { 36 | // return json.Marshal(p) 37 | // } 38 | // 39 | // func (p *ConnectRequst) Unmarshal(b []byte) error { 40 | // return json.Unmarshal(b, p) 41 | // } 42 | // 43 | // type ConnectReply struct { 44 | // Code int 45 | // Token string 46 | // } 47 | // 48 | // func (p *ConnectReply) Marshal() ([]byte, error) { 49 | // return json.Marshal(p) 50 | // } 51 | // 52 | // func (p *ConnectReply) Unmarshal(b []byte) error { 53 | // return json.Unmarshal(b, p) 54 | // } 55 | // 56 | // Server 57 | // type server struct { 58 | // *bdmsg.Server 59 | // } 60 | // 61 | // func newService(l net.Listener, handshakeTO time.Duration, 62 | // pumperInN, pumperOutN int) *server { 63 | // 64 | // s := &server{} 65 | // 66 | // mux := bdmsg.NewPumpMux(nil) 67 | // mux.HandleFunc(MsgTypeConnect, s.handleConnect) 68 | // 69 | // s.Server = bdmsg.NewServerF(l, bdmsg.DefaultIOC, handshakeTO, 70 | // mux, pumperInN, pumperOutN) 71 | // s.Server.Start() 72 | // return s 73 | // } 74 | // 75 | // func (s *server) handleConnect(ctx context.Context, 76 | // p *bdmsg.Pumper, t bdmsg.MsgType, m bdmsg.Msg) { 77 | // 78 | // msc := p.UserData().(*bdmsg.SClient) 79 | // if msc.Handshaked() { 80 | // panic(errors.New("Unexpected")) 81 | // } 82 | // 83 | // var request ConnectRequst 84 | // request.Unmarshal(m) // unmarshal request 85 | // 86 | // // process connect request 87 | // 88 | // // tell bdmsg that client is authorized 89 | // msc.Handshake() 90 | // 91 | // var reply ConnectReply 92 | // // init reply 93 | // mr, _ := reply.Marshal() // marshal reply 94 | // 95 | // msc.Output(MsgTypeConnectReply, mr) 96 | // } 97 | // Client 98 | // type client struct { 99 | // *bdmsg.Client 100 | // connected chan bool 101 | // } 102 | // 103 | // func newClient(conn net.Conn, pumperInN, pumperOutN int) *client { 104 | // 105 | // c := &client{connected: make(chan bool)} 106 | // 107 | // mux := bdmsg.NewPumpMux(nil) 108 | // mux.HandleFunc(MsgTypeConnectReply, c.handleConnectReply) 109 | // 110 | // c.Client = bdmsg.NewClient(nil, conn, bdmsg.DefaultIOC, 111 | // mux, pumperInN, pumperOutN) 112 | // c.doConnect() 113 | // return c 114 | // } 115 | // 116 | // func (c *client) doConnect() { 117 | // var request ConnectRequst 118 | // // init request 119 | // mr, _ := request.Marshal() // marshal request 120 | // 121 | // c.Client.Output(MsgTypeConnect, mr) 122 | // } 123 | // 124 | // func (c *client) handleConnectReply(ctx context.Context, 125 | // p *bdmsg.Pumper, t bdmsg.MsgType, m bdmsg.Msg) { 126 | // 127 | // var reply ConnectReply 128 | // reply.Unmarshal(m) // unmarshal reply 129 | // 130 | // // process connect reply 131 | // 132 | // close(c.connected) 133 | // } 134 | package bdmsg 135 | -------------------------------------------------------------------------------- /example_test.go: -------------------------------------------------------------------------------- 1 | // Copyright 2011 The Go Authors. All rights reserved. 2 | // Use of this source code is governed by a BSD-style 3 | // license that can be found in the LICENSE file. 4 | 5 | package bdmsg_test 6 | 7 | import ( 8 | "encoding/json" 9 | "errors" 10 | "github.com/someonegg/bdmsg" 11 | "golang.org/x/net/context" 12 | "net" 13 | "testing" 14 | "time" 15 | ) 16 | 17 | const ( 18 | MsgTypeConnect = 1 19 | MsgTypeConnectReply = 2 20 | ) 21 | 22 | type ConnectRequst struct { 23 | Name string 24 | Pass string 25 | } 26 | 27 | func (p *ConnectRequst) Marshal() ([]byte, error) { 28 | return json.Marshal(p) 29 | } 30 | 31 | func (p *ConnectRequst) Unmarshal(b []byte) error { 32 | return json.Unmarshal(b, p) 33 | } 34 | 35 | type ConnectReply struct { 36 | Code int 37 | Token string 38 | } 39 | 40 | func (p *ConnectReply) Marshal() ([]byte, error) { 41 | return json.Marshal(p) 42 | } 43 | 44 | func (p *ConnectReply) Unmarshal(b []byte) error { 45 | return json.Unmarshal(b, p) 46 | } 47 | 48 | type server struct { 49 | *bdmsg.Server 50 | } 51 | 52 | func newService(l net.Listener, handshakeTO time.Duration, 53 | pumperInN, pumperOutN int) *server { 54 | 55 | s := &server{} 56 | 57 | mux := bdmsg.NewPumpMux(nil) 58 | mux.HandleFunc(MsgTypeConnect, s.handleConnect) 59 | 60 | s.Server = bdmsg.NewServerF(l, bdmsg.DefaultIOC, handshakeTO, 61 | mux, pumperInN, pumperOutN) 62 | s.Server.Start() 63 | return s 64 | } 65 | 66 | func (s *server) handleConnect(ctx context.Context, 67 | p *bdmsg.Pumper, t bdmsg.MsgType, m bdmsg.Msg) { 68 | 69 | msc := p.UserData().(*bdmsg.SClient) 70 | if msc.Handshaked() { 71 | panic(errors.New("Unexpected")) 72 | } 73 | 74 | var request ConnectRequst 75 | request.Unmarshal(m) // unmarshal request 76 | 77 | // process connect request 78 | 79 | // tell bdmsg that client is authorized 80 | msc.Handshake() 81 | 82 | var reply ConnectReply 83 | // init reply 84 | mr, _ := reply.Marshal() // marshal reply 85 | 86 | msc.Output(MsgTypeConnectReply, mr) 87 | } 88 | 89 | type client struct { 90 | *bdmsg.Client 91 | connected chan bool 92 | } 93 | 94 | func newClient(conn net.Conn, pumperInN, pumperOutN int) *client { 95 | 96 | c := &client{connected: make(chan bool)} 97 | 98 | mux := bdmsg.NewPumpMux(nil) 99 | mux.HandleFunc(MsgTypeConnectReply, c.handleConnectReply) 100 | 101 | c.Client = bdmsg.NewClient(nil, conn, bdmsg.DefaultIOC, 102 | mux, pumperInN, pumperOutN) 103 | c.doConnect() 104 | return c 105 | } 106 | 107 | func (c *client) doConnect() { 108 | var request ConnectRequst 109 | // init request 110 | mr, _ := request.Marshal() // marshal request 111 | 112 | c.Client.Output(MsgTypeConnect, mr) 113 | } 114 | 115 | func (c *client) handleConnectReply(ctx context.Context, 116 | p *bdmsg.Pumper, t bdmsg.MsgType, m bdmsg.Msg) { 117 | 118 | var reply ConnectReply 119 | reply.Unmarshal(m) // unmarshal reply 120 | 121 | // process connect reply 122 | 123 | close(c.connected) 124 | } 125 | 126 | func TestExample(t *testing.T) { 127 | addr := "127.0.0.1:12345" 128 | 129 | l, err := net.Listen("tcp", addr) 130 | if err != nil { 131 | t.Fatal("listen failed: ", err) 132 | } 133 | svc := newService(l, time.Second, 10, 10) 134 | 135 | c, err := net.Dial("tcp", "127.0.0.1:12345") 136 | if err != nil { 137 | t.Fatal("dial failed: ", err) 138 | } 139 | cli := newClient(c, 10, 10) 140 | 141 | select { 142 | case <-cli.connected: 143 | case <-time.After(time.Second): 144 | t.Fatal("client connect failed") 145 | } 146 | 147 | svc.Stop() 148 | 149 | select { 150 | case <-svc.StopD(): 151 | case <-time.After(time.Second): 152 | t.Fatal("server stop failed") 153 | } 154 | 155 | select { 156 | case <-cli.StopD(): 157 | case <-time.After(time.Second): 158 | t.Fatal("client stop failed") 159 | } 160 | } 161 | -------------------------------------------------------------------------------- /io.go: -------------------------------------------------------------------------------- 1 | // Copyright 2015 someonegg. All rights reserved. 2 | // Use of this source code is governed by a BSD-style 3 | // license that can be found in the LICENSE file. 4 | 5 | package bdmsg 6 | 7 | import ( 8 | "bufio" 9 | "encoding/binary" 10 | "errors" 11 | "fmt" 12 | "github.com/someonegg/gocontainer/bufpool" 13 | "io" 14 | "sync" 15 | ) 16 | 17 | var ( 18 | ErrMsgTooBig = errors.New("message is too big") 19 | ErrMsgPacket = errors.New("message packet is wrong") 20 | ErrPackMsg = errors.New("pack message failed") 21 | ) 22 | 23 | // Message is a variable-length byte array. 24 | type Msg []byte 25 | 26 | // Message is distinguished by type. 27 | type MsgType int32 28 | 29 | type MsgReader interface { 30 | ReadMsg() (t MsgType, m Msg, err error) 31 | } 32 | 33 | type MsgWriter interface { 34 | WriteMsg(t MsgType, m Msg) (err error) 35 | } 36 | 37 | type MsgReadWriter interface { 38 | MsgReader 39 | MsgWriter 40 | } 41 | 42 | /* 43 | MsgRWIO implementes the MsgReadWriter interface. 44 | In the transport layer, message's layout is: 45 | 46 | A(length) + B(type) + C(data) 47 | A is 4-bytes int, big-endian 48 | B is 4-bytes int, big-endian 49 | C is byte array, its length is (A) - 4 50 | */ 51 | type MsgRWIO struct { 52 | RW io.ReadWriter 53 | F Flusher 54 | MsgMax int 55 | } 56 | 57 | type Flusher interface { 58 | Flush() error 59 | } 60 | 61 | func NewMsgRWIO(rw io.ReadWriter, msgMax int) *MsgRWIO { 62 | f, _ := rw.(Flusher) 63 | return &MsgRWIO{RW: rw, F: f, MsgMax: msgMax} 64 | } 65 | 66 | func (rw *MsgRWIO) ReadMsg() (t MsgType, m Msg, err error) { 67 | var _l int32 68 | err = binary.Read(rw.RW, binary.BigEndian, &_l) 69 | if err != nil { 70 | return 71 | } 72 | l := int(_l) 73 | 74 | if l < 4 { 75 | err = ErrMsgPacket 76 | return 77 | } 78 | l -= 4 79 | if l > rw.MsgMax { 80 | err = ErrMsgTooBig 81 | return 82 | } 83 | 84 | var _t int32 85 | err = binary.Read(rw.RW, binary.BigEndian, &_t) 86 | if err != nil { 87 | return 88 | } 89 | t = MsgType(_t) 90 | 91 | m = bufpool.Get(l) 92 | readed := 0 93 | for readed < l { 94 | var n int 95 | n, err = rw.RW.Read(m[readed:]) 96 | if n > 0 { 97 | readed += n 98 | continue 99 | } 100 | if err != nil { 101 | return 102 | } 103 | } 104 | 105 | return t, m, nil 106 | } 107 | 108 | func (rw *MsgRWIO) WriteMsg(t MsgType, m Msg) (err error) { 109 | l := len(m) 110 | if l > rw.MsgMax { 111 | err = ErrMsgTooBig 112 | return 113 | } 114 | l += 4 115 | 116 | err = binary.Write(rw.RW, binary.BigEndian, int32(l)) 117 | if err != nil { 118 | return 119 | } 120 | 121 | err = binary.Write(rw.RW, binary.BigEndian, int32(t)) 122 | if err != nil { 123 | return 124 | } 125 | 126 | n, err := rw.RW.Write(m) 127 | if err != nil { 128 | return 129 | } 130 | if (n + 4) != l { 131 | // Write must return a non-nil error if it returns n < len(p) 132 | err = ErrPackMsg 133 | return 134 | } 135 | 136 | if rw.F != nil { 137 | return rw.F.Flush() 138 | } 139 | return nil 140 | } 141 | 142 | type Converter interface { 143 | Convert(rw io.ReadWriter) MsgReadWriter 144 | } 145 | 146 | type StopNotifier interface { 147 | OnStop() 148 | } 149 | 150 | // DefaultMaxMsg is the default maximum message length. 151 | const DefaultMaxMsg = 128 * 1024 152 | 153 | type DefaultConverter struct { 154 | MsgMax int 155 | } 156 | 157 | func (c *DefaultConverter) Convert(rw io.ReadWriter) MsgReadWriter { 158 | return NewMsgRWIO(bufio.NewReadWriter( 159 | bufio.NewReader(rw), bufio.NewWriter(rw)), c.MsgMax) 160 | } 161 | 162 | // DefaultIOC is the default io.ReadWriter to MsgReadWriter converter. 163 | var DefaultIOC = Converter(&DefaultConverter{MsgMax: DefaultMaxMsg}) 164 | 165 | /* 166 | MsgRWDump provides message dump function. 167 | The dump format is: 168 | 169 | R|W\nMessageType\nMessageSize\nMessageData\n\n 170 | R for read, W for write 171 | MessageData part is raw data 172 | */ 173 | type MsgRWDump struct { 174 | rw MsgReadWriter 175 | careAbout func(MsgType, Msg) bool 176 | locker sync.Mutex 177 | dump io.ReadWriteCloser 178 | } 179 | 180 | func NewMsgRWDump(rw MsgReadWriter, 181 | careAbout func(MsgType, Msg) bool) *MsgRWDump { 182 | 183 | return &MsgRWDump{rw: rw, careAbout: careAbout} 184 | } 185 | 186 | func (rw *MsgRWDump) SetDump(dump io.ReadWriteCloser) io.ReadWriteCloser { 187 | rw.locker.Lock() 188 | defer rw.locker.Unlock() 189 | od := rw.dump 190 | rw.dump = dump 191 | return od 192 | } 193 | 194 | func (rw *MsgRWDump) Dump() io.ReadWriteCloser { 195 | rw.locker.Lock() 196 | defer rw.locker.Unlock() 197 | return rw.dump 198 | } 199 | 200 | func (rw *MsgRWDump) OnStop() { 201 | rw.locker.Lock() 202 | defer rw.locker.Unlock() 203 | if rw.dump != nil { 204 | rw.dump.Close() 205 | rw.dump = nil 206 | } 207 | } 208 | 209 | func (rw *MsgRWDump) needDump(t MsgType, m Msg) bool { 210 | if rw.careAbout != nil { 211 | return rw.careAbout(t, m) 212 | } 213 | return true 214 | } 215 | 216 | func (rw *MsgRWDump) ReadMsg() (t MsgType, m Msg, err error) { 217 | t, m, err = rw.rw.ReadMsg() 218 | if err != nil { 219 | return 220 | } 221 | 222 | if !rw.needDump(t, m) { 223 | return 224 | } 225 | 226 | rw.locker.Lock() 227 | defer rw.locker.Unlock() 228 | 229 | d := rw.dump 230 | if d == nil { 231 | return 232 | } 233 | 234 | fmt.Fprintf(d, "R %v %v\n", t, len(m)) 235 | d.Write(m) 236 | fmt.Fprintf(d, "\n\n") 237 | 238 | return 239 | } 240 | 241 | func (rw *MsgRWDump) WriteMsg(t MsgType, m Msg) (err error) { 242 | err = rw.rw.WriteMsg(t, m) 243 | if err != nil { 244 | return 245 | } 246 | 247 | if !rw.needDump(t, m) { 248 | return 249 | } 250 | 251 | rw.locker.Lock() 252 | defer rw.locker.Unlock() 253 | 254 | d := rw.dump 255 | if d == nil { 256 | return 257 | } 258 | 259 | fmt.Fprintf(d, "W %v %v\n", t, len(m)) 260 | d.Write(m) 261 | fmt.Fprintf(d, "\n\n") 262 | 263 | return 264 | } 265 | -------------------------------------------------------------------------------- /pumper.go: -------------------------------------------------------------------------------- 1 | // Copyright 2015 someonegg. All rights reserved. 2 | // Use of this source code is governed by a BSD-style 3 | // license that can be found in the LICENSE file. 4 | 5 | package bdmsg 6 | 7 | import ( 8 | "errors" 9 | "github.com/someonegg/gocontainer/bufpool" 10 | "github.com/someonegg/gox/syncx" 11 | "golang.org/x/net/context" 12 | "sync" 13 | "sync/atomic" 14 | ) 15 | 16 | var ( 17 | errUnknownPanic = errors.New("unknown panic") 18 | ) 19 | 20 | type msgEntry struct { 21 | t MsgType 22 | m Msg 23 | } 24 | 25 | // PumperHandler is the interface that can be registered 26 | // to process messages in the message pumper. 27 | type PumperHandler interface { 28 | Process(ctx context.Context, p *Pumper, t MsgType, m Msg) 29 | } 30 | 31 | type PumperStatis struct { 32 | // message count 33 | InTotal int64 34 | InProcess int64 35 | OutTotal int64 36 | OutProcess int64 37 | 38 | // message bytes 39 | BytesReaded int64 // from rw 40 | BytesWritten int64 // to rw 41 | } 42 | 43 | // Pumper represents a message pumper. It has a working loop 44 | // which reads, processes and writes messages continuously. 45 | // 46 | // Multiple goroutines can invoke methods on a Pumper simultaneously. 47 | type Pumper struct { 48 | err error 49 | quitF context.CancelFunc 50 | stopD syncx.DoneChan 51 | 52 | rw MsgReadWriter 53 | h PumperHandler 54 | ud interface{} 55 | 56 | // read 57 | rerr error 58 | rD syncx.DoneChan 59 | rQ chan msgEntry 60 | 61 | // write 62 | werr error 63 | wD syncx.DoneChan 64 | wQ chan msgEntry 65 | 66 | stat PumperStatis 67 | } 68 | 69 | // NewPumper allocates and returns a new Pumper. 70 | func NewPumper(rw MsgReadWriter, h PumperHandler, inN, outN int) *Pumper { 71 | t := &Pumper{} 72 | t.init(rw, h, inN, outN) 73 | return t 74 | } 75 | 76 | func (p *Pumper) init(rw MsgReadWriter, h PumperHandler, inN, outN int) { 77 | p.stopD = syncx.NewDoneChan() 78 | 79 | p.rw = rw 80 | p.h = h 81 | 82 | p.rD = syncx.NewDoneChan() 83 | p.rQ = make(chan msgEntry, inN) 84 | p.wD = syncx.NewDoneChan() 85 | p.wQ = make(chan msgEntry, outN) 86 | } 87 | 88 | // Start the pumper, parent and/or sn can be nil. 89 | // If sn is not nil, it will be called when the working loop exits. 90 | func (p *Pumper) Start(parent context.Context, sn StopNotifier) { 91 | if parent == nil { 92 | parent = context.Background() 93 | } 94 | 95 | var ctx context.Context 96 | ctx, p.quitF = context.WithCancel(parent) 97 | 98 | rwctx, rwqF := context.WithCancel(context.Background()) 99 | go p.reading(rwctx) 100 | go p.writing(rwctx) 101 | go p.work(ctx, rwqF, sn) 102 | } 103 | 104 | func (p *Pumper) work(ctx context.Context, 105 | rwqF context.CancelFunc, sn StopNotifier) { 106 | 107 | defer p.ending(rwqF, sn) 108 | 109 | for q := false; !q; { 110 | select { 111 | case <-ctx.Done(): 112 | q = true 113 | case e := <-p.rQ: 114 | atomic.AddInt64(&p.stat.InProcess, 1) 115 | p.procMsg(ctx, e.t, e.m) 116 | bufpool.Put(e.m) 117 | case <-p.rD: 118 | q = true 119 | case <-p.wD: 120 | q = true 121 | } 122 | } 123 | } 124 | 125 | func (p *Pumper) ending(rwqF context.CancelFunc, sn StopNotifier) { 126 | if e := recover(); e != nil { 127 | switch v := e.(type) { 128 | case error: 129 | p.err = v 130 | default: 131 | p.err = errUnknownPanic 132 | } 133 | } 134 | 135 | defer func() { recover() }() 136 | defer p.stopD.SetDone() 137 | 138 | // if ending from error. 139 | p.quitF() 140 | 141 | rwqF() 142 | if sn != nil { 143 | sn.OnStop() 144 | } 145 | 146 | <-p.rD 147 | <-p.wD 148 | } 149 | 150 | func (p *Pumper) procMsg(ctx context.Context, t MsgType, m Msg) { 151 | p.h.Process(ctx, p, t, m) 152 | } 153 | 154 | func (p *Pumper) reading(ctx context.Context) { 155 | defer func() { 156 | if e := recover(); e != nil { 157 | switch v := e.(type) { 158 | case error: 159 | p.rerr = v 160 | default: 161 | p.rerr = errUnknownPanic 162 | } 163 | } 164 | 165 | p.rD.SetDone() 166 | }() 167 | 168 | for q := false; !q; { 169 | t, m := p.readMsg() 170 | 171 | select { 172 | case <-ctx.Done(): 173 | q = true 174 | case p.rQ <- msgEntry{t, m}: 175 | atomic.AddInt64(&p.stat.InTotal, 1) 176 | } 177 | } 178 | } 179 | 180 | func (p *Pumper) readMsg() (MsgType, Msg) { 181 | t, m, err := p.rw.ReadMsg() 182 | if err != nil { 183 | panic(err) 184 | } 185 | atomic.AddInt64(&p.stat.BytesReaded, int64(len(m))) 186 | return t, m 187 | } 188 | 189 | func (p *Pumper) writing(ctx context.Context) { 190 | defer func() { 191 | if e := recover(); e != nil { 192 | switch v := e.(type) { 193 | case error: 194 | p.werr = v 195 | default: 196 | p.werr = errUnknownPanic 197 | } 198 | } 199 | 200 | p.wD.SetDone() 201 | }() 202 | 203 | for q := false; !q; { 204 | select { 205 | case <-ctx.Done(): 206 | q = true 207 | case e := <-p.wQ: 208 | atomic.AddInt64(&p.stat.OutProcess, 1) 209 | p.writeMsg(e.t, e.m) 210 | bufpool.Put(e.m) 211 | } 212 | } 213 | } 214 | 215 | func (p *Pumper) writeMsg(t MsgType, m Msg) { 216 | err := p.rw.WriteMsg(t, m) 217 | if err != nil { 218 | panic(err) 219 | } 220 | atomic.AddInt64(&p.stat.BytesWritten, int64(len(m))) 221 | } 222 | 223 | // Err returns non-nil if an error has happened. 224 | // When errored, the pumper will stop. 225 | func (p *Pumper) Err() error { 226 | if p.err != nil { 227 | return p.err 228 | } 229 | if p.rerr != nil { 230 | return p.rerr 231 | } 232 | return p.werr 233 | } 234 | 235 | // Stop requests to stop the working loop. 236 | func (p *Pumper) Stop() { 237 | p.quitF() 238 | } 239 | 240 | // StopD returns a done channel, it will be 241 | // signaled when the pumper is stopped. 242 | func (p *Pumper) StopD() syncx.DoneChanR { 243 | return p.stopD.R() 244 | } 245 | 246 | func (p *Pumper) Stopped() bool { 247 | return p.stopD.R().Done() 248 | } 249 | 250 | // Input copies the message data to the in-queue. 251 | func (p *Pumper) Input(t MsgType, m Msg) { 252 | cp := bufpool.Get(len(m)) 253 | copy(cp, m) 254 | select { 255 | case p.rQ <- msgEntry{t, cp}: 256 | atomic.AddInt64(&p.stat.InTotal, 1) 257 | case <-p.stopD: 258 | bufpool.Put(cp) 259 | } 260 | } 261 | 262 | // TryInput tries to copy the message data to the in-queue. 263 | func (p *Pumper) TryInput(t MsgType, m Msg) bool { 264 | cp := bufpool.Get(len(m)) 265 | copy(cp, m) 266 | select { 267 | case p.rQ <- msgEntry{t, cp}: 268 | atomic.AddInt64(&p.stat.InTotal, 1) 269 | return true 270 | default: 271 | bufpool.Put(cp) 272 | return false 273 | } 274 | } 275 | 276 | // Output copies the message data to the out-queue. 277 | func (p *Pumper) Output(t MsgType, m Msg) { 278 | cp := bufpool.Get(len(m)) 279 | copy(cp, m) 280 | select { 281 | case p.wQ <- msgEntry{t, cp}: 282 | atomic.AddInt64(&p.stat.OutTotal, 1) 283 | case <-p.stopD: 284 | bufpool.Put(cp) 285 | } 286 | } 287 | 288 | // TryOutput tries to copy the message data to the out-queue. 289 | func (p *Pumper) TryOutput(t MsgType, m Msg) bool { 290 | cp := bufpool.Get(len(m)) 291 | copy(cp, m) 292 | select { 293 | case p.wQ <- msgEntry{t, cp}: 294 | atomic.AddInt64(&p.stat.OutTotal, 1) 295 | return true 296 | default: 297 | bufpool.Put(cp) 298 | return false 299 | } 300 | } 301 | 302 | func (p *Pumper) Statis() *PumperStatis { 303 | return &PumperStatis{ 304 | InTotal: atomic.LoadInt64(&p.stat.InTotal), 305 | InProcess: atomic.LoadInt64(&p.stat.InProcess), 306 | OutTotal: atomic.LoadInt64(&p.stat.OutTotal), 307 | OutProcess: atomic.LoadInt64(&p.stat.OutProcess), 308 | 309 | BytesReaded: atomic.LoadInt64(&p.stat.BytesReaded), 310 | BytesWritten: atomic.LoadInt64(&p.stat.BytesWritten), 311 | } 312 | } 313 | 314 | // UserData returns the data user setted. 315 | func (p *Pumper) UserData() interface{} { 316 | return p.ud 317 | } 318 | 319 | func (p *Pumper) SetUserData(ud interface{}) { 320 | p.ud = ud 321 | } 322 | 323 | // InnerMsgRW returns the inner message readwriter. 324 | // You should make sure that its implements support concurrently 325 | // access if you want to call its methods. 326 | func (p *Pumper) InnerMsgRW() MsgReadWriter { 327 | return p.rw 328 | } 329 | 330 | // PumpMux is an message request multiplexer. 331 | // 332 | // It matches the type of each message against a list of registered 333 | // types and calls the matched handler. 334 | // 335 | // Multiple goroutines can invoke methods on a PumpMux simultaneously. 336 | type PumpMux struct { 337 | mu sync.RWMutex 338 | m map[MsgType]PumperHandler 339 | orphan PumperHandler 340 | } 341 | 342 | // NewPumpMux allocates and returns a new PumpMux. 343 | func NewPumpMux(orphan PumperHandler) *PumpMux { 344 | return &PumpMux{ 345 | m: make(map[MsgType]PumperHandler), 346 | orphan: orphan, 347 | } 348 | } 349 | 350 | // Process dispatches the request to the handler whose 351 | // type matches the message type. 352 | func (mux *PumpMux) Process(ctx context.Context, p *Pumper, t MsgType, m Msg) { 353 | mux.mu.RLock() 354 | defer mux.mu.RUnlock() 355 | 356 | if h, ok := mux.m[t]; ok { 357 | h.Process(ctx, p, t, m) 358 | } else { 359 | if mux.orphan != nil { 360 | mux.orphan.Process(ctx, p, t, m) 361 | } 362 | } 363 | } 364 | 365 | // Handle registers the handler for the given type. 366 | // If a handler already exists for type, Handle panics. 367 | func (mux *PumpMux) Handle(t MsgType, h PumperHandler) { 368 | mux.mu.Lock() 369 | defer mux.mu.Unlock() 370 | 371 | if h == nil { 372 | panic("msgpump: nil handler") 373 | } 374 | if _, ok := mux.m[t]; ok { 375 | panic("msgpump: multiple registrations") 376 | } 377 | 378 | mux.m[t] = h 379 | } 380 | 381 | // HandleFunc registers the handler function for the given type. 382 | func (mux *PumpMux) HandleFunc(t MsgType, h func(context.Context, 383 | *Pumper, MsgType, Msg)) { 384 | 385 | mux.Handle(t, HandlerFunc(h)) 386 | } 387 | 388 | // The HandlerFunc type is an adapter to allow the use of 389 | // ordinary functions as message handlers. If f is a function 390 | // with the appropriate signature, HandlerFunc(f) is a 391 | // PumperHandler object that calls f. 392 | type HandlerFunc func(context.Context, *Pumper, MsgType, Msg) 393 | 394 | // Process calls f(ctx, p, t, m). 395 | func (f HandlerFunc) Process(ctx context.Context, p *Pumper, t MsgType, m Msg) { 396 | f(ctx, p, t, m) 397 | } 398 | -------------------------------------------------------------------------------- /server.go: -------------------------------------------------------------------------------- 1 | // Copyright 2015 someonegg. All rights reserved. 2 | // Use of this source code is governed by a BSD-style 3 | // license that can be found in the LICENSE file. 4 | 5 | package bdmsg 6 | 7 | import ( 8 | "github.com/someonegg/gox/syncx" 9 | "golang.org/x/net/context" 10 | "net" 11 | "sync" 12 | "sync/atomic" 13 | "time" 14 | ) 15 | 16 | // SClient represents a message client in the server side. 17 | // You can send(recive) message to(from) client use SClient.Pumper(). 18 | // 19 | // Multiple goroutines can invoke methods on a SClient simultaneously. 20 | type SClient struct { 21 | Pumper 22 | 23 | c net.Conn 24 | wg *sync.WaitGroup 25 | 26 | handshaked int32 27 | } 28 | 29 | // NewSClient allocates and returns a new SClient. 30 | // 31 | // The ownership of c will be transferred to SClient, dont 32 | // control it in other places. 33 | func NewSClient(parent context.Context, c net.Conn, ioc Converter, 34 | h PumperHandler, pumperInN, pumperOutN int, wg *sync.WaitGroup) *SClient { 35 | 36 | rw := ioc.Convert(c) 37 | 38 | t := &SClient{} 39 | t.c = c 40 | t.wg = wg 41 | t.Pumper.init(rw, h, pumperInN, pumperOutN) 42 | t.Pumper.SetUserData(t) 43 | t.Pumper.Start(parent, t) 44 | 45 | return t 46 | } 47 | 48 | func (c *SClient) OnStop() { 49 | c.c.Close() 50 | if sn, ok := c.rw.(StopNotifier); ok { 51 | sn.OnStop() 52 | } 53 | if c.wg != nil { 54 | c.wg.Done() 55 | } 56 | } 57 | 58 | // InnerPumper is c.Pumper, its initial userdata is *SClient. 59 | func (c *SClient) InnerPumper() *Pumper { 60 | return &c.Pumper 61 | } 62 | 63 | func (c *SClient) Conn() net.Conn { 64 | return c.c 65 | } 66 | 67 | // Handshake tell bdmsg that client is authorized. 68 | // Unauthorized client will be closed after HSTO. 69 | func (c *SClient) Handshake() { 70 | atomic.StoreInt32(&c.handshaked, 1) 71 | } 72 | 73 | func (c *SClient) Handshaked() bool { 74 | return atomic.LoadInt32(&c.handshaked) == 1 75 | } 76 | 77 | // Server represents a message server. 78 | // You can accept message clients use it. 79 | // 80 | // Multiple goroutines can invoke methods on a Server simultaneously. 81 | type Server struct { 82 | err error 83 | quitF context.CancelFunc 84 | stopD syncx.DoneChan 85 | 86 | l net.Listener 87 | ioc Converter 88 | hsto time.Duration 89 | 90 | h PumperHandler 91 | pumperInN int 92 | pumperOutN int 93 | 94 | cliWG sync.WaitGroup 95 | } 96 | 97 | // NewServerF allocates and returns a new Server. 98 | // 99 | // Note: hsto is "handshake timeout". 100 | func NewServerF(l net.Listener, ioc Converter, hsto time.Duration, 101 | h PumperHandler, pumperInN, pumperOutN int) *Server { 102 | 103 | s := &Server{ 104 | stopD: syncx.NewDoneChan(), 105 | l: l, 106 | ioc: ioc, 107 | hsto: hsto, 108 | h: h, 109 | pumperInN: pumperInN, 110 | pumperOutN: pumperOutN, 111 | } 112 | 113 | return s 114 | } 115 | 116 | func (s *Server) Start() { 117 | ctx, quitF := context.WithCancel(context.Background()) 118 | s.quitF = quitF 119 | go s.work(ctx) 120 | } 121 | 122 | func (s *Server) work(ctx context.Context) { 123 | defer s.ending() 124 | 125 | var tempDelay time.Duration // how long to sleep on accept failure 126 | 127 | for q := false; !q; { 128 | c, err := s.l.Accept() 129 | if err != nil { 130 | if nerr, ok := err.(net.Error); ok && nerr.Temporary() { 131 | if tempDelay == 0 { 132 | tempDelay = 5 * time.Millisecond 133 | } else { 134 | tempDelay *= 2 135 | } 136 | if max := 1 * time.Second; tempDelay > max { 137 | tempDelay = max 138 | } 139 | } else { 140 | panic(err) 141 | } 142 | } else { 143 | tempDelay = 0 144 | } 145 | 146 | if c != nil { 147 | s.newClient(ctx, c) 148 | } 149 | 150 | select { 151 | case <-ctx.Done(): 152 | q = true 153 | continue 154 | default: 155 | } 156 | 157 | if tempDelay > 0 { 158 | time.Sleep(tempDelay) 159 | } 160 | } 161 | } 162 | 163 | func (s *Server) ending() { 164 | if e := recover(); e != nil { 165 | switch v := e.(type) { 166 | case error: 167 | s.err = v 168 | default: 169 | s.err = errUnknownPanic 170 | } 171 | } 172 | 173 | s.quitF() 174 | s.stopD.SetDone() 175 | } 176 | 177 | func (s *Server) newClient(ctx context.Context, c net.Conn) { 178 | s.cliWG.Add(1) 179 | cli := NewSClient(ctx, c, s.ioc, 180 | s.h, s.pumperInN, s.pumperOutN, &s.cliWG) 181 | 182 | if s.hsto > 0 { 183 | go monitorHSTO(cli, s.hsto) 184 | } 185 | } 186 | 187 | func monitorHSTO(c *SClient, hsto time.Duration) { 188 | defer func() { recover() }() 189 | select { 190 | case <-time.After(hsto): 191 | if !c.Handshaked() { 192 | c.Stop() 193 | } 194 | } 195 | } 196 | 197 | // Err returns non-nil if an error has happened. 198 | // When errored, the server will stop. 199 | func (s *Server) Err() error { 200 | return s.err 201 | } 202 | 203 | // Stop requests to stop the server. 204 | func (s *Server) Stop() { 205 | s.quitF() 206 | s.l.Close() 207 | } 208 | 209 | // StopD returns a done channel, it will be 210 | // signaled when the server is stopped. 211 | func (s *Server) StopD() syncx.DoneChanR { 212 | return s.stopD.R() 213 | } 214 | 215 | func (s *Server) Stopped() bool { 216 | return s.stopD.R().Done() 217 | } 218 | 219 | func (s *Server) WaitClients() { 220 | s.cliWG.Wait() 221 | } 222 | --------------------------------------------------------------------------------