micro/server/rpc_server.go
2015-05-23 17:40:53 +01:00

147 lines
2.5 KiB
Go

package server
import (
"bytes"
"sync"
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"
)
type rpcServer struct {
mtx sync.RWMutex
address string
opts options
rpc *rpc.Server
exit chan chan error
}
func newRpcServer(address string, opt ...Option) Server {
var opts options
for _, o := range opt {
o(&opts)
}
if opts.transport == nil {
opts.transport = transport.DefaultTransport
}
return &rpcServer{
opts: opts,
address: address,
rpc: rpc.NewServer(),
exit: make(chan chan error),
}
}
func (s *rpcServer) accept(sock transport.Socket) {
var msg transport.Message
if err := sock.Recv(&msg); err != nil {
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
}
// 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
}
sock.Send(&transport.Message{
Header: map[string]string{
"Content-Type": ct,
},
Body: rsp.Bytes(),
})
}
func (s *rpcServer) Address() string {
s.mtx.RLock()
address := s.address
s.mtx.RUnlock()
return address
}
func (s *rpcServer) Init() error {
return nil
}
func (s *rpcServer) NewReceiver(handler interface{}) Receiver {
return newRpcReceiver("", handler)
}
func (s *rpcServer) NewNamedReceiver(name string, handler interface{}) Receiver {
return newRpcReceiver(name, handler)
}
func (s *rpcServer) Register(r Receiver) error {
if len(r.Name()) > 0 {
s.rpc.RegisterName(r.Name(), r.Handler())
return nil
}
s.rpc.Register(r.Handler())
return nil
}
func (s *rpcServer) Start() error {
registerHealthChecker(s)
ts, err := s.opts.transport.Listen(s.address)
if err != nil {
return err
}
log.Infof("Listening on %s", ts.Addr())
s.mtx.RLock()
s.address = ts.Addr()
s.mtx.RUnlock()
go ts.Accept(s.accept)
go func() {
ch := <-s.exit
ch <- ts.Close()
}()
return nil
}
func (s *rpcServer) Stop() error {
ch := make(chan error)
s.exit <- ch
return <-ch
}