mirror of
https://github.com/go-micro/go-micro.git
synced 2024-12-24 10:07:04 +02:00
Fix rpc go routine leak
This commit is contained in:
parent
266b6dbc64
commit
af94899b54
@ -81,6 +81,20 @@ type router struct {
|
||||
subscribers map[string][]*subscriber
|
||||
}
|
||||
|
||||
// rpcRouter encapsulates functions that become a server.Router
|
||||
type rpcRouter struct {
|
||||
h func(context.Context, Request, interface{}) error
|
||||
m func(context.Context, Message) error
|
||||
}
|
||||
|
||||
func (r rpcRouter) ProcessMessage(ctx context.Context, msg Message) error {
|
||||
return r.m(ctx, msg)
|
||||
}
|
||||
|
||||
func (r rpcRouter) ServeRequest(ctx context.Context, req Request, rsp Response) error {
|
||||
return r.h(ctx, req, rsp)
|
||||
}
|
||||
|
||||
func newRpcRouter() *router {
|
||||
return &router{
|
||||
serviceMap: make(map[string]*service),
|
||||
|
@ -58,19 +58,6 @@ func newRpcServer(opts ...Option) Server {
|
||||
}
|
||||
}
|
||||
|
||||
type rpcRouter struct {
|
||||
h func(context.Context, Request, interface{}) error
|
||||
m func(context.Context, Message) error
|
||||
}
|
||||
|
||||
func (r rpcRouter) ProcessMessage(ctx context.Context, msg Message) error {
|
||||
return r.m(ctx, msg)
|
||||
}
|
||||
|
||||
func (r rpcRouter) ServeRequest(ctx context.Context, req Request, rsp Response) error {
|
||||
return r.h(ctx, req, rsp)
|
||||
}
|
||||
|
||||
// HandleEvent handles inbound messages to the service directly
|
||||
// TODO: handle requests from an event. We won't send a response.
|
||||
func (s *rpcServer) HandleEvent(e broker.Event) error {
|
||||
@ -141,26 +128,34 @@ func (s *rpcServer) HandleEvent(e broker.Event) error {
|
||||
|
||||
// ServeConn serves a single connection
|
||||
func (s *rpcServer) ServeConn(sock transport.Socket) {
|
||||
var wg sync.WaitGroup
|
||||
var mtx sync.RWMutex
|
||||
// global error tracking
|
||||
var gerr error
|
||||
// streams are multiplexed on Micro-Stream or Micro-Id header
|
||||
sockets := make(map[string]*socket.Socket)
|
||||
pool := socket.NewPool()
|
||||
|
||||
// get global waitgroup
|
||||
s.Lock()
|
||||
gg := s.wg
|
||||
s.Unlock()
|
||||
|
||||
// waitgroup to wait for processing to finish
|
||||
wg := &waitGroup{
|
||||
gg: gg,
|
||||
}
|
||||
|
||||
defer func() {
|
||||
// only wait if there's no error
|
||||
if gerr == nil {
|
||||
// wait till done
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
// close all the sockets for this connection
|
||||
pool.Close()
|
||||
|
||||
// close underlying socket
|
||||
sock.Close()
|
||||
|
||||
// close the sockets
|
||||
mtx.Lock()
|
||||
for id, psock := range sockets {
|
||||
psock.Close()
|
||||
delete(sockets, id)
|
||||
}
|
||||
mtx.Unlock()
|
||||
|
||||
// recover any panics
|
||||
if r := recover(); r != nil {
|
||||
log.Log("panic recovered: ", r)
|
||||
@ -170,7 +165,12 @@ func (s *rpcServer) ServeConn(sock transport.Socket) {
|
||||
|
||||
for {
|
||||
var msg transport.Message
|
||||
// process inbound messages one at a time
|
||||
if err := sock.Recv(&msg); err != nil {
|
||||
// set a global error and return
|
||||
// we're saying we essentially can't
|
||||
// use the socket anymore
|
||||
gerr = err
|
||||
return
|
||||
}
|
||||
|
||||
@ -185,9 +185,12 @@ func (s *rpcServer) ServeConn(sock transport.Socket) {
|
||||
msg.Header["Micro-Error"] = err.Error()
|
||||
}
|
||||
// write back some 200
|
||||
sock.Send(&transport.Message{
|
||||
if err := sock.Send(&transport.Message{
|
||||
Header: msg.Header,
|
||||
})
|
||||
}); err != nil {
|
||||
gerr = err
|
||||
break
|
||||
}
|
||||
// we're done
|
||||
continue
|
||||
}
|
||||
@ -205,57 +208,52 @@ func (s *rpcServer) ServeConn(sock transport.Socket) {
|
||||
id = msg.Header["Micro-Id"]
|
||||
}
|
||||
|
||||
// check stream id
|
||||
var stream bool
|
||||
|
||||
if v := getHeader("Micro-Stream", msg.Header); len(v) > 0 {
|
||||
stream = true
|
||||
}
|
||||
|
||||
// check if we have an existing socket
|
||||
psock, ok := pool.Get(id)
|
||||
|
||||
// if we don't have a socket and its a stream
|
||||
if !ok && stream {
|
||||
// check if its a last stream EOS error
|
||||
err := msg.Header["Micro-Error"]
|
||||
if err == lastStreamResponseError.Error() {
|
||||
pool.Release(psock)
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
// got an existing socket already
|
||||
if ok {
|
||||
// we're starting processing
|
||||
wg.Add(1)
|
||||
|
||||
// add to wait group if "wait" is opt-in
|
||||
s.Lock()
|
||||
swg := s.wg
|
||||
s.Unlock()
|
||||
|
||||
if swg != nil {
|
||||
swg.Add(1)
|
||||
}
|
||||
|
||||
// check we have an existing socket
|
||||
mtx.RLock()
|
||||
psock, ok := sockets[id]
|
||||
mtx.RUnlock()
|
||||
|
||||
// got the socket
|
||||
if ok {
|
||||
// accept the message
|
||||
// pass the message to that existing socket
|
||||
if err := psock.Accept(&msg); err != nil {
|
||||
// delete the socket
|
||||
mtx.Lock()
|
||||
delete(sockets, id)
|
||||
mtx.Unlock()
|
||||
}
|
||||
|
||||
// done(1)
|
||||
if swg != nil {
|
||||
swg.Done()
|
||||
// release the socket if there's an error
|
||||
pool.Release(psock)
|
||||
}
|
||||
|
||||
// done waiting
|
||||
wg.Done()
|
||||
|
||||
// continue to the next message
|
||||
continue
|
||||
}
|
||||
|
||||
// no socket was found
|
||||
psock = socket.New()
|
||||
// no socket was found so its new
|
||||
// set the local and remote values
|
||||
psock.SetLocal(sock.Local())
|
||||
psock.SetRemote(sock.Remote())
|
||||
|
||||
// load the socket
|
||||
// load the socket with the current message
|
||||
psock.Accept(&msg)
|
||||
|
||||
// save a new socket
|
||||
mtx.Lock()
|
||||
sockets[id] = psock
|
||||
mtx.Unlock()
|
||||
|
||||
// now walk the usual path
|
||||
|
||||
// we use this Timeout header to set a server deadline
|
||||
@ -292,38 +290,33 @@ func (s *rpcServer) ServeConn(sock transport.Socket) {
|
||||
// setup old protocol
|
||||
cf := setupProtocol(&msg)
|
||||
|
||||
// no old codec
|
||||
// no legacy codec needed
|
||||
if cf == nil {
|
||||
// TODO: needs better error handling
|
||||
var err error
|
||||
// try get a new codec
|
||||
if cf, err = s.newCodec(ct); err != nil {
|
||||
sock.Send(&transport.Message{
|
||||
// no codec found so send back an error
|
||||
if err := sock.Send(&transport.Message{
|
||||
Header: map[string]string{
|
||||
"Content-Type": "text/plain",
|
||||
},
|
||||
Body: []byte(err.Error()),
|
||||
})
|
||||
|
||||
if swg != nil {
|
||||
swg.Done()
|
||||
}); err != nil {
|
||||
gerr = err
|
||||
}
|
||||
|
||||
wg.Done()
|
||||
|
||||
return
|
||||
// release the socket we just created
|
||||
pool.Release(psock)
|
||||
// now continue
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
// create a new rpc codec based on the pseudo socket and codec
|
||||
rcodec := newRpcCodec(&msg, psock, cf)
|
||||
// check the protocol as well
|
||||
protocol := rcodec.String()
|
||||
|
||||
// check stream id
|
||||
var stream bool
|
||||
|
||||
if v := getHeader("Micro-Stream", msg.Header); len(v) > 0 {
|
||||
stream = true
|
||||
}
|
||||
|
||||
// internal request
|
||||
request := &rpcRequest{
|
||||
service: getHeader("Micro-Service", msg.Header),
|
||||
@ -363,11 +356,11 @@ func (s *rpcServer) ServeConn(sock transport.Socket) {
|
||||
r = rpcRouter{h: handler}
|
||||
}
|
||||
|
||||
// process the outbound messages from the socket
|
||||
go func(id string, psock *socket.Socket) {
|
||||
// wait for processing to exit
|
||||
wg.Add(1)
|
||||
|
||||
// process the outbound messages from the socket
|
||||
go func(id string, psock *socket.Socket) {
|
||||
defer func() {
|
||||
// TODO: don't hack this but if its grpc just break out of the stream
|
||||
// We do this because the underlying connection is h2 and its a stream
|
||||
@ -375,7 +368,9 @@ func (s *rpcServer) ServeConn(sock transport.Socket) {
|
||||
case "grpc":
|
||||
sock.Close()
|
||||
}
|
||||
|
||||
// release the socket
|
||||
pool.Release(psock)
|
||||
// signal we're done
|
||||
wg.Done()
|
||||
}()
|
||||
|
||||
@ -383,10 +378,6 @@ func (s *rpcServer) ServeConn(sock transport.Socket) {
|
||||
// get the message from our internal handler/stream
|
||||
m := new(transport.Message)
|
||||
if err := psock.Process(m); err != nil {
|
||||
// delete the socket
|
||||
mtx.Lock()
|
||||
delete(sockets, id)
|
||||
mtx.Unlock()
|
||||
return
|
||||
}
|
||||
|
||||
@ -399,7 +390,15 @@ func (s *rpcServer) ServeConn(sock transport.Socket) {
|
||||
|
||||
// serve the request in a go routine as this may be a stream
|
||||
go func(id string, psock *socket.Socket) {
|
||||
defer psock.Close()
|
||||
// add to the waitgroup
|
||||
wg.Add(1)
|
||||
|
||||
defer func() {
|
||||
// release the socket
|
||||
pool.Release(psock)
|
||||
// signal we're done
|
||||
wg.Done()
|
||||
}()
|
||||
|
||||
// serve the actual request using the request router
|
||||
if serveRequestError := r.ServeRequest(ctx, request, response); serveRequestError != nil {
|
||||
@ -416,21 +415,9 @@ func (s *rpcServer) ServeConn(sock transport.Socket) {
|
||||
|
||||
// could not write error response
|
||||
if writeError != nil && !alreadyClosed {
|
||||
log.Logf("rpc: unable to write error response: %v", writeError)
|
||||
log.Debugf("rpc: unable to write error response: %v", writeError)
|
||||
}
|
||||
}
|
||||
|
||||
mtx.Lock()
|
||||
delete(sockets, id)
|
||||
mtx.Unlock()
|
||||
|
||||
// signal we're done
|
||||
if swg != nil {
|
||||
swg.Done()
|
||||
}
|
||||
|
||||
// done with this socket
|
||||
wg.Done()
|
||||
}(id, psock)
|
||||
}
|
||||
}
|
||||
|
32
server/rpc_util.go
Normal file
32
server/rpc_util.go
Normal file
@ -0,0 +1,32 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"sync"
|
||||
)
|
||||
|
||||
// waitgroup for global management of connections
|
||||
type waitGroup struct {
|
||||
// local waitgroup
|
||||
lg sync.WaitGroup
|
||||
// global waitgroup
|
||||
gg *sync.WaitGroup
|
||||
}
|
||||
|
||||
func (w *waitGroup) Add(i int) {
|
||||
w.lg.Add(i)
|
||||
if w.gg != nil {
|
||||
w.gg.Add(i)
|
||||
}
|
||||
}
|
||||
|
||||
func (w *waitGroup) Done() {
|
||||
w.lg.Done()
|
||||
if w.gg != nil {
|
||||
w.gg.Done()
|
||||
}
|
||||
}
|
||||
|
||||
func (w *waitGroup) Wait() {
|
||||
// only wait on local group
|
||||
w.lg.Wait()
|
||||
}
|
@ -285,8 +285,14 @@ func (h *httpTransportSocket) Recv(m *Message) error {
|
||||
|
||||
func (h *httpTransportSocket) Send(m *Message) error {
|
||||
if h.r.ProtoMajor == 1 {
|
||||
// make copy of header
|
||||
hdr := make(http.Header)
|
||||
for k, v := range h.r.Header {
|
||||
hdr[k] = v
|
||||
}
|
||||
|
||||
rsp := &http.Response{
|
||||
Header: h.r.Header,
|
||||
Header: hdr,
|
||||
Body: ioutil.NopCloser(bytes.NewReader(m.Body)),
|
||||
Status: "200 OK",
|
||||
StatusCode: 200,
|
||||
|
56
util/socket/pool.go
Normal file
56
util/socket/pool.go
Normal file
@ -0,0 +1,56 @@
|
||||
package socket
|
||||
|
||||
import (
|
||||
"sync"
|
||||
)
|
||||
|
||||
type Pool struct {
|
||||
sync.RWMutex
|
||||
pool map[string]*Socket
|
||||
}
|
||||
|
||||
func (p *Pool) Get(id string) (*Socket, bool) {
|
||||
// attempt to get existing socket
|
||||
p.RLock()
|
||||
socket, ok := p.pool[id]
|
||||
if ok {
|
||||
p.RUnlock()
|
||||
return socket, ok
|
||||
}
|
||||
p.RUnlock()
|
||||
|
||||
// create new socket
|
||||
socket = New(id)
|
||||
// save socket
|
||||
p.Lock()
|
||||
p.pool[id] = socket
|
||||
p.Unlock()
|
||||
// return socket
|
||||
return socket, false
|
||||
}
|
||||
|
||||
func (p *Pool) Release(s *Socket) {
|
||||
p.Lock()
|
||||
defer p.Unlock()
|
||||
|
||||
// close the socket
|
||||
s.Close()
|
||||
delete(p.pool, s.id)
|
||||
}
|
||||
|
||||
// Close the pool and delete all the sockets
|
||||
func (p *Pool) Close() {
|
||||
p.Lock()
|
||||
defer p.Unlock()
|
||||
for id, sock := range p.pool {
|
||||
sock.Close()
|
||||
delete(p.pool, id)
|
||||
}
|
||||
}
|
||||
|
||||
// NewPool returns a new socket pool
|
||||
func NewPool() *Pool {
|
||||
return &Pool{
|
||||
pool: make(map[string]*Socket),
|
||||
}
|
||||
}
|
@ -9,6 +9,7 @@ import (
|
||||
|
||||
// Socket is our pseudo socket for transport.Socket
|
||||
type Socket struct {
|
||||
id string
|
||||
// closed
|
||||
closed chan bool
|
||||
// remote addr
|
||||
@ -119,8 +120,9 @@ func (s *Socket) Close() error {
|
||||
// New returns a new pseudo socket which can be used in the place of a transport socket.
|
||||
// Messages are sent to the socket via Accept and receives from the socket via Process.
|
||||
// SetLocal/SetRemote should be called before using the socket.
|
||||
func New() *Socket {
|
||||
func New(id string) *Socket {
|
||||
return &Socket{
|
||||
id: id,
|
||||
closed: make(chan bool),
|
||||
local: "local",
|
||||
remote: "remote",
|
||||
|
Loading…
Reference in New Issue
Block a user