k2safe/go-omni
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: bound native shutdown and close connections that become idle during drain

Browse Source
This commit is contained in:
k2safe
2026-09-23 12:57:27 +08:00
parent 8e18528fc2
commit be620d2243
3 changed files with 179 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
transport/native/server.go
View File

@@ -41,6 +41,7 @@ const (
// It supports keep-alive, graceful shutdown, bounded concurrency and standard
// http.Handler values. Zero MaxConnections and MaxConcurrentRequests mean
// unlimited; pooled protocol objects are reused automatically.
// A Server owns one serving lifecycle; create a new Server to restart.
type Server struct {
ReadHeaderTimeout time.Duration
ReadTimeout time.Duration
@@ -232,9 +233,9 @@ func (s *Server) Serve(ctx context.Context, listener net.Listener) error {
}
s.mu.Lock()
if s.listener != nil {
if s.cancel != nil {
s.mu.Unlock()
return errors.New("go-omni/native: server is already running")
return errors.New("go-omni/native: server has already been started")
}
s.listener = listener
s.closing = false
@@ -357,7 +358,8 @@ func (s *Server) Shutdown(ctx context.Context) error {
for _, state := range states {
state.close()
}
<-done
// Handlers must cooperate with cancellation. Waiting here would make
// an uncooperative handler defeat the caller's shutdown deadline.
return ctx.Err()
}
}
@@ -446,8 +448,7 @@ func (s *Server) serveConnection(baseCtx context.Context, state *connection) {
cancel()
writer.reset("")
s.writerPool.Put(writer)
state.active.Store(false)
s.activeRequests.Add(-1)
s.finishRequest(state)
return
default:
s.rejectedRequests.Add(1)
@@ -459,8 +460,7 @@ func (s *Server) serveConnection(baseCtx context.Context, state *connection) {
cancel()
writer.reset("")
s.writerPool.Put(writer)
state.active.Store(false)
s.activeRequests.Add(-1)
closing = s.finishRequest(state) || closing
if closeAfter || closing {
return
}
@@ -496,8 +496,7 @@ func (s *Server) serveConnection(baseCtx context.Context, state *connection) {
}
writer.reset("")
s.writerPool.Put(writer)
state.active.Store(false)
s.activeRequests.Add(-1)
s.finishRequest(state)
return
}
// Request bodies are streamed into the handler, so a malformed
@@ -518,8 +517,7 @@ func (s *Server) serveConnection(baseCtx context.Context, state *connection) {
}
writer.reset("")
s.writerPool.Put(writer)
state.active.Store(false)
s.activeRequests.Add(-1)
s.finishRequest(state)
return
}
if requestSlot {
@@ -528,14 +526,25 @@ func (s *Server) serveConnection(baseCtx context.Context, state *connection) {
shouldClose := closeAfter || writer.closeAfter || closing
writer.reset("")
s.writerPool.Put(writer)
state.active.Store(false)
s.activeRequests.Add(-1)
shouldClose = s.finishRequest(state) || shouldClose
if shouldClose {
return
}
}
}
// Publish the idle transition under the same lock that starts shutdown. Either
// Shutdown observes an idle connection and closes it, or this request observes
// shutdown and closes it instead of starting another Keep-Alive read.
func (s *Server) finishRequest(state *connection) bool {
s.mu.Lock()
state.active.Store(false)
closing := s.closing || state.closeWhenIdle.Load()
s.mu.Unlock()
s.activeRequests.Add(-1)
return closing
}
func (c *connection) close() {
if c.closed.CompareAndSwap(false, true) {
defer close(c.closeDone)