add support for streaming requests. cleanup watcher initilisation

server/buffer.go

@@ -5,8 +5,7 @@ import (
 )
 
 type buffer struct {
-	io.Reader
-	io.Writer
+	io.ReadWriter
 }
 
 func (b *buffer) Close() error {

server/options.go

@@ -1,10 +1,12 @@
 package server
 
 import (
+	"github.com/myodc/go-micro/registry"
 	"github.com/myodc/go-micro/transport"
 )
 
 type options struct {
+	registry  registry.Registry
 	transport transport.Transport
 	metadata  map[string]string
 	name      string
@@ -19,6 +21,10 @@ func newOptions(opt ...Option) options {
 		o(&opts)
 	}
 
+	if opts.registry == nil {
+		opts.registry = registry.DefaultRegistry
+	}
+
 	if opts.transport == nil {
 		opts.transport = transport.DefaultTransport
 	}

server/rpc_codec.go

@@ -0,0 +1,82 @@
+package server
+
+import (
+	"bytes"
+	"fmt"
+
+	"github.com/myodc/go-micro/transport"
+	rpc "github.com/youtube/vitess/go/rpcplus"
+	js "github.com/youtube/vitess/go/rpcplus/jsonrpc"
+	pb "github.com/youtube/vitess/go/rpcplus/pbrpc"
+)
+
+type rpcPlusCodec struct {
+	socket transport.Socket
+	codec  rpc.ServerCodec
+
+	req *transport.Message
+
+	wbuf *bytes.Buffer
+	rbuf *bytes.Buffer
+}
+
+func newRpcPlusCodec(req *transport.Message, socket transport.Socket) *rpcPlusCodec {
+	return &rpcPlusCodec{
+		socket: socket,
+		req:    req,
+		wbuf:   bytes.NewBuffer(nil),
+		rbuf:   bytes.NewBuffer(nil),
+	}
+}
+
+func (c *rpcPlusCodec) ReadRequestHeader(r *rpc.Request) error {
+	c.rbuf.Reset()
+	c.rbuf.Write(c.req.Body)
+	buf := &buffer{c.rbuf}
+
+	switch c.req.Header["Content-Type"] {
+	case "application/octet-stream":
+		c.codec = pb.NewServerCodec(buf)
+	case "application/json":
+		c.codec = js.NewServerCodec(buf)
+	default:
+		return fmt.Errorf("unsupported content type %s", c.req.Header["Content-Type"])
+	}
+
+	return c.codec.ReadRequestHeader(r)
+}
+
+func (c *rpcPlusCodec) ReadRequestBody(r interface{}) error {
+	return c.codec.ReadRequestBody(r)
+}
+
+func (c *rpcPlusCodec) WriteResponse(r *rpc.Response, body interface{}, last bool) error {
+	c.wbuf.Reset()
+	buf := &buffer{c.wbuf}
+
+	var cc rpc.ServerCodec
+	switch c.req.Header["Content-Type"] {
+	case "application/octet-stream":
+		cc = pb.NewServerCodec(buf)
+	case "application/json":
+		cc = js.NewServerCodec(buf)
+	default:
+		return fmt.Errorf("unsupported request type: %s", c.req.Header["Content-Type"])
+	}
+
+	if err := cc.WriteResponse(r, body, last); err != nil {
+		return err
+	}
+
+	return c.socket.Send(&transport.Message{
+		Header: map[string]string{"Content-Type": c.req.Header["Content-Type"]},
+		Body:   c.wbuf.Bytes(),
+	})
+
+}
+
+func (c *rpcPlusCodec) Close() error {
+	c.wbuf.Reset()
+	c.rbuf.Reset()
+	return c.socket.Close()
+}

server/rpc_server.go

@@ -1,15 +1,11 @@
 package server
 
 import (
-	"bytes"
-
 	c "github.com/myodc/go-micro/context"
 	"github.com/myodc/go-micro/transport"
 
 	log "github.com/golang/glog"
 	rpc "github.com/youtube/vitess/go/rpcplus"
-	js "github.com/youtube/vitess/go/rpcplus/jsonrpc"
-	pb "github.com/youtube/vitess/go/rpcplus/pbrpc"
 
 	"golang.org/x/net/context"
 )
@@ -34,42 +30,17 @@ func (s *rpcServer) accept(sock transport.Socket) {
 		return
 	}
 
-	rbq := bytes.NewBuffer(msg.Body)
-	rsp := bytes.NewBuffer(nil)
-	defer rsp.Reset()
-	defer rbq.Reset()
-
-	buf := &buffer{
-		rbq,
-		rsp,
-	}
-
-	var cc rpc.ServerCodec
-	switch msg.Header["Content-Type"] {
-	case "application/octet-stream":
-		cc = pb.NewServerCodec(buf)
-	case "application/json":
-		cc = js.NewServerCodec(buf)
-	default:
-		return
-	}
+	codec := newRpcPlusCodec(&msg, sock)
 
 	// strip our headers
-	ct := msg.Header["Content-Type"]
-	delete(msg.Header, "Content-Type")
-
-	ctx := c.WithMetadata(context.Background(), msg.Header)
-
-	if err := s.rpc.ServeRequestWithContext(ctx, cc); err != nil {
-		return
+	hdr := make(map[string]string)
+	for k, v := range msg.Header {
+		hdr[k] = v
 	}
+	delete(hdr, "Content-Type")
 
-	sock.Send(&transport.Message{
-		Header: map[string]string{
-			"Content-Type": ct,
-		},
-		Body: rsp.Bytes(),
-	})
+	ctx := c.WithMetadata(context.Background(), hdr)
+	s.rpc.ServeRequestWithContext(ctx, codec)
 }
 
 func (s *rpcServer) Config() options {

server/server.go

@@ -50,6 +50,12 @@ func Address(a string) Option {
 	}
 }
 
+func Registry(r registry.Registry) Option {
+	return func(o *options) {
+		o.registry = r
+	}
+}
+
 func Transport(t transport.Transport) Option {
 	return func(o *options) {
 		o.transport = t
@@ -121,9 +127,9 @@ func Run() error {
 
 	log.Infof("Registering node: %s", node.Id)
 
-	err := registry.Register(service)
+	err := config.registry.Register(service)
 	if err != nil {
-		log.Fatal("Failed to register: %v", err)
+		log.Fatalf("Failed to register: %v", err)
 	}
 
 	ch := make(chan os.Signal, 1)
@@ -131,7 +137,7 @@ func Run() error {
 	log.Infof("Received signal %s", <-ch)
 
 	log.Infof("Deregistering %s", node.Id)
-	registry.Deregister(service)
+	config.registry.Deregister(service)
 
 	return Stop()
 }