diff --git a/agent/grpc/handler.go b/agent/grpc/handler.go index d70fd2b10c..e78231276d 100644 --- a/agent/grpc/handler.go +++ b/agent/grpc/handler.go @@ -22,7 +22,7 @@ func NewHandler(addr net.Addr, register func(server *grpc.Server)) *Handler { ) register(srv) - lis := &chanListener{addr: addr, conns: make(chan net.Conn)} + lis := &chanListener{addr: addr, conns: make(chan net.Conn), done: make(chan struct{})} return &Handler{srv: srv, listener: lis} } @@ -51,22 +51,22 @@ func (h *Handler) Shutdown() error { type chanListener struct { conns chan net.Conn addr net.Addr + done chan struct{} } // Accept blocks until a connection is received from Handle, and then returns the // connection. Accept implements part of the net.Listener interface for grpc.Server. func (l *chanListener) Accept() (net.Conn, error) { select { - case c, ok := <-l.conns: - if !ok { - return nil, &net.OpError{ - Op: "accept", - Net: l.addr.Network(), - Addr: l.addr, - Err: fmt.Errorf("listener closed"), - } - } + case c := <-l.conns: return c, nil + case <-l.done: + return nil, &net.OpError{ + Op: "accept", + Net: l.addr.Network(), + Addr: l.addr, + Err: fmt.Errorf("listener closed"), + } } } @@ -75,7 +75,7 @@ func (l *chanListener) Addr() net.Addr { } func (l *chanListener) Close() error { - close(l.conns) + close(l.done) return nil }