micro/gateway/rpc/server.go

package rpc

import (
	"errors"
	"net"
	"sync"
)

var (
	ErrHandleExists = errors.New("handle already exists")
)

type HandleFunc func(ctx *Context) error

type Server struct {
	listener   net.Listener
	ch         chan *Request
	ctxPool    sync.Pool
	serviceMap sync.Map // map[string]HandleFunc
	exitChan   chan struct{}
}

func (svr *Server) getContext() *Context {
	if v := svr.ctxPool.Get(); v != nil {
		return v.(*Context)
	} else {
		return &Context{}
	}
}

func (svr *Server) putContext(c *Context) {
	svr.ctxPool.Put(c)
}

func (svr *Server) wrkLoop() {
	for {
		select {
		case req, ok := <-svr.ch:
			if ok {
				svr.handleRequest(req)
			}
		case <-svr.exitChan:
			return
		}
	}
}

func (svr *Server) handleRequest(req *Request) {
	var (
		ok  bool
		err error
		cb  HandleFunc
		val interface{}
		ctx *Context
	)
	if val, ok = svr.serviceMap.Load(req); ok {
		cb = val.(HandleFunc)
		ctx = svr.getContext()
		resp := &Response{}
		ctx.Reset(req, resp)
		if err = cb(ctx); err != nil {
			_ = writeFrame(req.conn, &Frame{
				Func: FuncResponse,
				Data: resp.Encode(),
			})
		}
	}
}

func (svr *Server) process(conn net.Conn) {
	var (
		err   error
		frame *Frame
	)
	defer func() {
		_ = conn.Close()
	}()
	for {
		if frame, err = readFrame(conn); err != nil {
			return
		}
		switch frame.Func {
		case FuncPing:
			_ = writeFrame(conn, &Frame{Func: FuncPing})
		case FuncRequest:
			req := ReadRequest(frame.Data)
			req.Reset(conn)
			select {
			case svr.ch <- req:
			default:
			}
		}
	}
}

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)
	}
	return
}

func (svr *Server) Serve(l net.Listener) {
	for {
		if conn, err := l.Accept(); err == nil {
			go svr.process(conn)
		} else {
			break
		}
	}
}

func (svr *Server) Close() (err error) {
	err = svr.listener.Close()
	return
}

func NewServer() *Server {
	return &Server{}
}