checkpoint

server/rpc_stream.go

@@ -0,0 +1,79 @@
+package server
+
+import (
+	"errors"
+	"io"
+	"log"
+	"sync"
+
+	"golang.org/x/net/context"
+)
+
+type rpcStream struct {
+	sync.RWMutex
+	seq     uint64
+	closed  bool
+	err     error
+	request Request
+	codec   serverCodec
+	context context.Context
+}
+
+func (r *rpcStream) Context() context.Context {
+	return r.context
+}
+
+func (r *rpcStream) Request() Request {
+	return r.request
+}
+
+func (r *rpcStream) Send(msg interface{}) error {
+	r.Lock()
+	defer r.Unlock()
+
+	seq := r.seq
+	r.seq++
+
+	resp := response{
+		ServiceMethod: r.request.Method(),
+		Seq:           seq,
+	}
+
+	err := codec.WriteResponse(&resp, msg, false)
+	if err != nil {
+		log.Println("rpc: writing response:", err)
+	}
+	return err
+}
+
+func (r *rpcStream) Recv(msg interface{}) error {
+	r.Lock()
+	defer r.Unlock()
+
+	req := request{}
+
+	if err := codec.ReadRequestHeader(&req); err != nil {
+		// discard body
+		codec.ReadRequestBody(nil)
+		return err
+	}
+
+	if err = codec.ReadRequestBody(msg); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (r *rpcStream) Error() error {
+	r.RLock()
+	defer r.RUnlock()
+	return r.err
+}
+
+func (r *rpcStream) Close() error {
+	r.Lock()
+	defer r.Unlock()
+	r.closed = true
+	return r.codec.Close()
+}