rpc.go 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152
  1. // Copyright (c) 2012, 2013 Ugorji Nwoke. All rights reserved.
  2. // Use of this source code is governed by a BSD-style license found in the LICENSE file.
  3. package codec
  4. import (
  5. "bufio"
  6. "io"
  7. "net/rpc"
  8. "sync"
  9. )
  10. // Rpc provides a rpc Server or Client Codec for rpc communication.
  11. type Rpc interface {
  12. ServerCodec(conn io.ReadWriteCloser, h Handle) rpc.ServerCodec
  13. ClientCodec(conn io.ReadWriteCloser, h Handle) rpc.ClientCodec
  14. }
  15. // RpcCodecBuffered allows access to the underlying bufio.Reader/Writer
  16. // used by the rpc connection. It accomodates use-cases where the connection
  17. // should be used by rpc and non-rpc functions, e.g. streaming a file after
  18. // sending an rpc response.
  19. type RpcCodecBuffered interface {
  20. BufferedReader() *bufio.Reader
  21. BufferedWriter() *bufio.Writer
  22. }
  23. // -------------------------------------
  24. // rpcCodec defines the struct members and common methods.
  25. type rpcCodec struct {
  26. rwc io.ReadWriteCloser
  27. dec *Decoder
  28. enc *Encoder
  29. bw *bufio.Writer
  30. br *bufio.Reader
  31. mu sync.Mutex
  32. cls bool
  33. }
  34. func newRPCCodec(conn io.ReadWriteCloser, h Handle) rpcCodec {
  35. bw := bufio.NewWriter(conn)
  36. br := bufio.NewReader(conn)
  37. return rpcCodec{
  38. rwc: conn,
  39. bw: bw,
  40. br: br,
  41. enc: NewEncoder(bw, h),
  42. dec: NewDecoder(br, h),
  43. }
  44. }
  45. func (c *rpcCodec) BufferedReader() *bufio.Reader {
  46. return c.br
  47. }
  48. func (c *rpcCodec) BufferedWriter() *bufio.Writer {
  49. return c.bw
  50. }
  51. func (c *rpcCodec) write(obj1, obj2 interface{}, writeObj2, doFlush bool) (err error) {
  52. if c.cls {
  53. return io.EOF
  54. }
  55. if err = c.enc.Encode(obj1); err != nil {
  56. return
  57. }
  58. if writeObj2 {
  59. if err = c.enc.Encode(obj2); err != nil {
  60. return
  61. }
  62. }
  63. if doFlush && c.bw != nil {
  64. return c.bw.Flush()
  65. }
  66. return
  67. }
  68. func (c *rpcCodec) read(obj interface{}) (err error) {
  69. if c.cls {
  70. return io.EOF
  71. }
  72. //If nil is passed in, we should still attempt to read content to nowhere.
  73. if obj == nil {
  74. var obj2 interface{}
  75. return c.dec.Decode(&obj2)
  76. }
  77. return c.dec.Decode(obj)
  78. }
  79. func (c *rpcCodec) Close() error {
  80. if c.cls {
  81. return io.EOF
  82. }
  83. c.cls = true
  84. return c.rwc.Close()
  85. }
  86. func (c *rpcCodec) ReadResponseBody(body interface{}) error {
  87. return c.read(body)
  88. }
  89. // -------------------------------------
  90. type goRpcCodec struct {
  91. rpcCodec
  92. }
  93. func (c *goRpcCodec) WriteRequest(r *rpc.Request, body interface{}) error {
  94. // Must protect for concurrent access as per API
  95. c.mu.Lock()
  96. defer c.mu.Unlock()
  97. return c.write(r, body, true, true)
  98. }
  99. func (c *goRpcCodec) WriteResponse(r *rpc.Response, body interface{}) error {
  100. c.mu.Lock()
  101. defer c.mu.Unlock()
  102. return c.write(r, body, true, true)
  103. }
  104. func (c *goRpcCodec) ReadResponseHeader(r *rpc.Response) error {
  105. return c.read(r)
  106. }
  107. func (c *goRpcCodec) ReadRequestHeader(r *rpc.Request) error {
  108. return c.read(r)
  109. }
  110. func (c *goRpcCodec) ReadRequestBody(body interface{}) error {
  111. return c.read(body)
  112. }
  113. // -------------------------------------
  114. // goRpc is the implementation of Rpc that uses the communication protocol
  115. // as defined in net/rpc package.
  116. type goRpc struct{}
  117. // GoRpc implements Rpc using the communication protocol defined in net/rpc package.
  118. // Its methods (ServerCodec and ClientCodec) return values that implement RpcCodecBuffered.
  119. var GoRpc goRpc
  120. func (x goRpc) ServerCodec(conn io.ReadWriteCloser, h Handle) rpc.ServerCodec {
  121. return &goRpcCodec{newRPCCodec(conn, h)}
  122. }
  123. func (x goRpc) ClientCodec(conn io.ReadWriteCloser, h Handle) rpc.ClientCodec {
  124. return &goRpcCodec{newRPCCodec(conn, h)}
  125. }
  126. var _ RpcCodecBuffered = (*rpcCodec)(nil) // ensure *rpcCodec implements RpcCodecBuffered