|
@@ -1,13 +1,10 @@
|
|
|
package rpc
|
|
|
|
|
|
import (
|
|
|
- "errors"
|
|
|
+ "git.nspix.com/golang/micro/log"
|
|
|
"net"
|
|
|
"sync"
|
|
|
-)
|
|
|
-
|
|
|
-var (
|
|
|
- ErrHandleExists = errors.New("handle already exists")
|
|
|
+ "sync/atomic"
|
|
|
)
|
|
|
|
|
|
type HandleFunc func(ctx *Context) error
|
|
@@ -17,6 +14,7 @@ type Server struct {
|
|
|
ch chan *Request
|
|
|
ctxPool sync.Pool
|
|
|
serviceMap sync.Map // map[string]HandleFunc
|
|
|
+ exitFlag int32
|
|
|
exitChan chan struct{}
|
|
|
}
|
|
|
|
|
@@ -53,16 +51,36 @@ func (svr *Server) handleRequest(req *Request) {
|
|
|
val interface{}
|
|
|
ctx *Context
|
|
|
)
|
|
|
- if val, ok = svr.serviceMap.Load(req); ok {
|
|
|
+ log.Debugf("rcp handle request %s@%d", req.Method, req.Sequence)
|
|
|
+ if val, ok = svr.serviceMap.Load(req.Method); ok {
|
|
|
cb = val.(HandleFunc)
|
|
|
ctx = svr.getContext()
|
|
|
- resp := &Response{}
|
|
|
+ ctx.Reset(req, NewResponse())
|
|
|
+ if err = cb(ctx); err == nil {
|
|
|
+ if err = writeFrame(req.conn, &Frame{
|
|
|
+ Func: FuncResponse,
|
|
|
+ Sequence: req.Sequence,
|
|
|
+ Data: ctx.Response().Bytes(),
|
|
|
+ }); err != nil {
|
|
|
+ log.Warnf("rcp write %d response error: %s", req.Sequence, err.Error())
|
|
|
+ } else {
|
|
|
+ log.Debugf("rcp write %d response success", req.Sequence)
|
|
|
+ }
|
|
|
+ }
|
|
|
+ } else {
|
|
|
+ ctx = svr.getContext()
|
|
|
+ resp := NewResponse()
|
|
|
+ resp.code = 404
|
|
|
+ resp.message = "not found"
|
|
|
ctx.Reset(req, resp)
|
|
|
- if err = cb(ctx); err != nil {
|
|
|
- _ = writeFrame(req.conn, &Frame{
|
|
|
- Func: FuncResponse,
|
|
|
- Data: resp.Encode(),
|
|
|
- })
|
|
|
+ if err = writeFrame(req.conn, &Frame{
|
|
|
+ Func: FuncResponse,
|
|
|
+ Sequence: req.Sequence,
|
|
|
+ Data: ctx.Response().Bytes(),
|
|
|
+ }); err != nil {
|
|
|
+ log.Warnf("rcp write %d response error: %s", req.Sequence, err.Error())
|
|
|
+ } else {
|
|
|
+ log.Debugf("rcp write %d response success", req.Sequence)
|
|
|
}
|
|
|
}
|
|
|
}
|
|
@@ -81,43 +99,57 @@ func (svr *Server) process(conn net.Conn) {
|
|
|
}
|
|
|
switch frame.Func {
|
|
|
case FuncPing:
|
|
|
- _ = writeFrame(conn, &Frame{Func: FuncPing})
|
|
|
+ if err = writeFrame(conn, &Frame{Func: FuncPing}); err != nil {
|
|
|
+ log.Warnf("rcp write ping frame error: %s", err.Error())
|
|
|
+ }
|
|
|
case FuncRequest:
|
|
|
- req := ReadRequest(frame.Data)
|
|
|
- req.Reset(conn)
|
|
|
- select {
|
|
|
- case svr.ch <- req:
|
|
|
- default:
|
|
|
+ //读取一个请求
|
|
|
+ if req, err2 := ReadRequest(frame.Data); err2 == nil {
|
|
|
+ req.reset(frame.Sequence, conn)
|
|
|
+ log.Debugf("rcp read request %s@%d", req.Method, req.Sequence)
|
|
|
+ select {
|
|
|
+ case svr.ch <- req:
|
|
|
+ default:
|
|
|
+ }
|
|
|
+ } else {
|
|
|
+ log.Warnf("read rpc request error: %s", err2.Error())
|
|
|
}
|
|
|
}
|
|
|
}
|
|
|
}
|
|
|
|
|
|
-func (svr *Server) HandleFunc(method string, f HandleFunc) (err error) {
|
|
|
- _, ok := svr.serviceMap.Load(method)
|
|
|
- if ok {
|
|
|
- err = ErrHandleExists
|
|
|
- } else {
|
|
|
- svr.serviceMap.Store(method, f)
|
|
|
- }
|
|
|
+func (svr *Server) Handle(method string, f HandleFunc) {
|
|
|
+ svr.serviceMap.Store(method, f)
|
|
|
return
|
|
|
}
|
|
|
|
|
|
-func (svr *Server) Serve(l net.Listener) {
|
|
|
+func (svr *Server) Serve(l net.Listener) (err error) {
|
|
|
+ svr.listener = l
|
|
|
+ go func() {
|
|
|
+ svr.wrkLoop()
|
|
|
+ }()
|
|
|
for {
|
|
|
- if conn, err := l.Accept(); err == nil {
|
|
|
+ if conn, err2 := svr.listener.Accept(); err2 == nil {
|
|
|
go svr.process(conn)
|
|
|
} else {
|
|
|
+ err = err2
|
|
|
break
|
|
|
}
|
|
|
}
|
|
|
+ return
|
|
|
}
|
|
|
|
|
|
func (svr *Server) Close() (err error) {
|
|
|
- err = svr.listener.Close()
|
|
|
+ if atomic.CompareAndSwapInt32(&svr.exitFlag, 0, 1) {
|
|
|
+ err = svr.listener.Close()
|
|
|
+ close(svr.exitChan)
|
|
|
+ }
|
|
|
return
|
|
|
}
|
|
|
|
|
|
func NewServer() *Server {
|
|
|
- return &Server{}
|
|
|
+ return &Server{
|
|
|
+ ch: make(chan *Request, 10),
|
|
|
+ exitChan: make(chan struct{}),
|
|
|
+ }
|
|
|
}
|