Hijacked connection will not be handled by the graceful shutdown

Need to fix this

pkg/proxy/server.go

@@ -177,8 +177,9 @@ type server struct {
  donec chan struct{}
  errc chan error
 
- closeOnce sync.Once
- closeWg sync.WaitGroup
+ closeOnce sync.Once
+ closeWg sync.WaitGroup
+ closeHijackedConn sync.WaitGroup
 
  listenerMu sync.RWMutex
  listener *customListener
@@ -279,7 +280,6 @@ func NewServer(cfg ServerConfig) Server {
  }
 
  s.closeWg.Add(1)
-
  var ln net.Listener
  var err error
  if !s.tlsInfo.Empty() {
@@ -332,7 +332,8 @@ func (c customListener) Accept() (net.Conn, error) {
  c.s.pauseAcceptMu.Unlock()
  select {
  case <-pausec:
- default:
+ case <-c.s.donec:
+ return nil, nil
  }
 
  c.s.latencyAcceptMu.RLock()
@@ -341,24 +342,44 @@ func (c customListener) Accept() (net.Conn, error) {
  if lat > 0 {
  select {
  case <-time.After(lat):
- default:
+ case <-c.s.donec:
+ return nil, nil
  }
  }
 
  conn, err := c.l.Accept()
  if err != nil {
- c.s.errc <- err
+ select {
+ case c.s.errc <- err:
+ select {
+ case <-c.s.donec:
+ return nil, nil
+ default:
+ }
+ case <-c.s.donec:
+ return nil, nil
+ }
  c.s.lg.Debug("listener accept error", zap.Error(err))
 
  if strings.HasSuffix(err.Error(), "use of closed network connection") {
  select {
  case <-time.After(c.s.retryInterval):
- default:
+ case <-c.s.donec:
+ return nil, nil
  }
  c.s.lg.Debug("listener is closed; retry listening on", zap.String("from", c.s.From()))
 
  if err = c.s.ResetListener(); err != nil {
- c.s.errc <- err
+ select {
+ case c.s.errc <- err:
+ select {
+ case <-c.s.donec:
+ return nil, nil
+ default:
+ }
+ case <-c.s.donec:
+ return nil, nil
+ }
  c.s.lg.Warn("failed to reset listener", zap.Error(err))
  }
  }
@@ -386,7 +407,16 @@ func (s *serverHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
  hijacker, _ := resp.(http.Hijacker)
  in, _, err := hijacker.Hijack()
  if err != nil {
- s.s.errc <- err
+ select {
+ case s.s.errc <- err:
+ select {
+ case <-s.s.donec:
+ return
+ default:
+ }
+ case <-s.s.donec:
+ return
+ }
  s.s.lg.Debug("ServeHTTP hijack error", zap.Error(err))
  panic(err)
  }
@@ -407,7 +437,16 @@ func (s *serverHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
  var tp *http.Transport
  tp, err = transport.NewTransport(s.s.tlsInfo, s.s.dialTimeout)
  if err != nil {
- s.s.errc <- err
+ select {
+ case s.s.errc <- err:
+ select {
+ case <-s.s.donec:
+ return
+ default:
+ }
+ case <-s.s.donec:
+ return
+ }
  s.s.lg.Debug("failed to get new Transport", zap.Error(err))
  return
  }
@@ -416,29 +455,47 @@ func (s *serverHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
  out, err = net.Dial(targetScheme, targetHost)
  }
  if err != nil {
- s.s.errc <- err
+ select {
+ case s.s.errc <- err:
+ select {
+ case <-s.s.donec:
+ return
+ default:
+ }
+ case <-s.s.donec:
+ return
+ }
  s.s.lg.Debug("failed to dial", zap.Error(err))
  return
  }
 
  var dstPort int
  dstPort, err = getPort(out.RemoteAddr())
  if err != nil {
- s.s.errc <- err
+ select {
+ case s.s.errc <- err:
+ select {
+ case <-s.s.donec:
+ return
+ default:
+ }
+ case <-s.s.donec:
+ return
+ }
  s.s.lg.Debug("failed to parse port in transmit", zap.Error(err))
  return
  }
 
- s.s.closeWg.Add(2)
+ s.s.closeHijackedConn.Add(2)
  go func() {
- defer s.s.closeWg.Done()
+ defer s.s.closeHijackedConn.Done()
  // read incoming bytes from listener, dispatch to outgoing connection
  s.s.transmit(out, in, dstPort)
  out.Close()
  in.Close()
  }()
  go func() {
- defer s.s.closeWg.Done()
+ defer s.s.closeHijackedConn.Done()
  // read response from outgoing connection, write back to listener
  s.s.receive(in, out, dstPort)
  in.Close()
@@ -710,27 +767,28 @@ func (s *server) Close() (err error) {
  s.closeOnce.Do(func() {
  close(s.donec)
 
- if s.httpServer != nil {
- if err = s.httpServer.Shutdown(context.TODO()); err != nil {
- return
- }
- s.httpServer = nil
- } else {
- s.listenerMu.Lock()
-
- if s.listener != nil {
- err = s.listener.Close()
- s.lg.Info(
- "closed proxy listener",
- zap.String("from", s.From()),
- zap.String("to", s.To()),
- )
- }
- s.lg.Sync()
- s.listenerMu.Unlock()
+ s.closeHijackedConn.Wait()
+
+ if err = s.httpServer.Shutdown(context.TODO()); err != nil {
+ return
  }
+ s.httpServer = nil
+
+ // s.listenerMu.Lock()
+ // if s.listener != nil {
+ // err = s.listener.Close()
+ // s.lg.Info(
+ // "closed proxy listener",
+ // zap.String("from", s.From()),
+ // zap.String("to", s.To()),
+ // )
+ // }
+ // s.lg.Sync()
+ // s.listenerMu.Unlock()
  })
- s.closeWg.Wait()
+
+ // s.closeWg.Wait()
+
  return err
 }
 

tests/e2e/blackhole_test.go

@@ -61,8 +61,8 @@ func blackholeTestByMockingPartition(t *testing.T, clusterSize int, partitionLea
  t.Logf("Blackholing traffic from and to member %q", partitionedMember.Config().Name)
  epc.BlackholePeer(partitionedMember)
 
- t.Logf("Wait 5s for any open connections to expire")
- time.Sleep(5 * time.Second)
+ t.Logf("Wait 1s for any open connections to expire")
+ time.Sleep(1 * time.Second)
 
  t.Logf("Wait for new leader election with remaining members")
  leaderEPC := epc.Procs[waitLeader(t, epc, mockPartitionNodeIndex)]
@@ -80,7 +80,7 @@ func blackholeTestByMockingPartition(t *testing.T, clusterSize int, partitionLea
  epc.UnblackholePeer(partitionedMember)
 
  leaderEPC = epc.Procs[epc.WaitLeader(t)]
- time.Sleep(5 * time.Second)
+ time.Sleep(1 * time.Second)
  assertRevision(t, leaderEPC, 21)
  assertRevision(t, partitionedMember, 21)
 }