ing

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)