diff --git a/servlet.go b/servlet.go index c113e66..c5d3c19 100644 --- a/servlet.go +++ b/servlet.go @@ -18,7 +18,7 @@ func NewServlet(sh ServletHandler) Servlet { } type Servlet interface { - Start(parentCTX cuc.Context, conn interface{}) error + Start(parentCTX cuc.Context, conn interface{}, doneChan chan<- error) error Stop() Send(method string, args ...interface{}) (err error) @@ -31,6 +31,7 @@ type servlet struct { sh ServletHandler messageQueueChan chan *messageState + doneChan chan<- error conn interface{} serverCodec protocol.ServerCodec @@ -38,7 +39,7 @@ type servlet struct { stopWg sync.WaitGroup } -func (s *servlet) Start(parentCTX cuc.Context, conn interface{}) error { +func (s *servlet) Start(parentCTX cuc.Context, conn interface{}, doneChan chan<- error) error { if nil == s.sh { panic("Servlet: servlet handler must be specified.") } @@ -54,6 +55,7 @@ func (s *servlet) Start(parentCTX cuc.Context, conn interface{}) error { return err } + s.doneChan = doneChan s.conn = conn s.serverCodec = sc @@ -102,7 +104,12 @@ func (s *servlet) Context() ServletContext { } func handleServlet(s *servlet) { - defer s.stopWg.Done() + var err error + + defer func() { + s.doneChan <- err + s.stopWg.Done() + }() subStopChan := make(chan struct{}) @@ -112,8 +119,6 @@ func handleServlet(s *servlet) { writerDone := make(chan error, 1) go handleWriter(s, subStopChan, writerDone) - var err error - select { case err = <-readerDone: close(subStopChan)