diff --git a/server/option.go b/server/option.go index 9c8a23c..016f31e 100644 --- a/server/option.go +++ b/server/option.go @@ -47,6 +47,13 @@ func WithPool(maxWorkers, maxCapacity int, options ...pond.Option) OptionFn { } } +// WithCustomPool uses a custom goroutine pool. +func WithCustomPool(pool WorkerPool) OptionFn { + return func(s *Server) { + s.pool = pool + } +} + // WithAsyncWrite sets AsyncWrite to true. func WithAsyncWrite() OptionFn { return func(s *Server) { diff --git a/server/server.go b/server/server.go index 79a093e..af299f7 100644 --- a/server/server.go +++ b/server/server.go @@ -20,7 +20,6 @@ import ( "sync/atomic" "time" - "github.com/alitto/pond" "github.com/smallnest/rpcx/log" "github.com/smallnest/rpcx/protocol" "github.com/smallnest/rpcx/share" @@ -69,6 +68,13 @@ var ( type Handler func(ctx *Context) error +type WorkerPool interface { + Submit(task func()) + StopAndWaitFor(deadline time.Duration) + Stop() + StopAndWait() +} + // Server is rpcx server that use TCP or UDP. type Server struct { ln net.Listener @@ -79,7 +85,7 @@ type Server struct { DisableHTTPGateway bool // disable http invoke or not. DisableJSONRPC bool // disable json rpc or not. AsyncWrite bool // set true if your server only serves few clients - pool *pond.WorkerPool + pool WorkerPool serviceMapMu sync.RWMutex serviceMap map[string]*service