diff --git a/src/xhttp/README.md b/src/xhttp/README.md index 2a97dc5..fdfcd40 100644 --- a/src/xhttp/README.md +++ b/src/xhttp/README.md @@ -25,6 +25,7 @@ go get github.com/mix-go/xhttp | xhttp.WithMiddlewares(middlewares ...Middleware) RequestOption | Set configuration item | | xhttp.BuildJSON(v interface{}) Body | Generate json string | | xhttp.BuildQuery(m map[string]string) Body | Generate urlencoded query string | +| xhttp.Shutdown(ctx context.Context) | Do shutdown | ## Debug Log @@ -51,10 +52,10 @@ The log object contains the following fields ```go type Log struct { - Duration time.Duration `json:"duration"` - Request *Request `json:"request"` // The Request.RetryAttempts field records the number of retry attempts - Response *Response `json:"response"` // If request error this field is equal to nil - Error error `json:"error"` + Duration time.Duration `json:"duration"` + Request *Request `json:"request"` // The Request.RetryAttempts field records the number of retry attempts + Response *Response `json:"response"` // If request error this field is equal to nil + Error error `json:"error"` } ``` @@ -123,7 +124,7 @@ ch := make(chan os.Signal) signal.Notify(ch, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM) go func() { <-ch - xhttp.Shutdown() + xhttp.Shutdown(context.Background()) }() ``` diff --git a/src/xhttp/shutdownctrl.go b/src/xhttp/shutdownctrl.go index ea50346..f235fcd 100644 --- a/src/xhttp/shutdownctrl.go +++ b/src/xhttp/shutdownctrl.go @@ -1,6 +1,7 @@ package xhttp import ( + "context" "sync" "sync/atomic" ) @@ -28,11 +29,22 @@ func (sc *ShutdownController) EndRequest() { sc.waitGroup.Done() } -func (sc *ShutdownController) InitiateShutdown() { +func (sc *ShutdownController) InitiateShutdown(ctx context.Context) { atomic.StoreInt32(&sc.shutdownFlag, 1) - sc.waitGroup.Wait() + + done := make(chan struct{}) + go func() { + sc.waitGroup.Wait() + close(done) + }() + select { + case <-ctx.Done(): // Timeout or ctx canceled, stop waiting. + return + case <-done: // waitGroup has completed + return + } } -func Shutdown() { - shutdownController.InitiateShutdown() +func Shutdown(ctx context.Context) { + shutdownController.InitiateShutdown(ctx) }