89 lines
1.8 KiB
Go
89 lines
1.8 KiB
Go
package app
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"sync"
|
|
"time"
|
|
|
|
"git.yuntuo.hk/k2safe/go-omni/lifecycle"
|
|
)
|
|
|
|
// Admission and completion share a lock: once stopped, no request can enter
|
|
// between observing zero active handlers and releasing their dependencies.
|
|
type requestDrain struct {
|
|
mu sync.Mutex
|
|
stopped bool
|
|
active map[*requestWork]struct{}
|
|
done chan struct{}
|
|
}
|
|
|
|
type requestWork struct{ cancel context.CancelFunc }
|
|
|
|
func (d *requestDrain) begin(ctx context.Context) (context.Context, *requestWork) {
|
|
d.mu.Lock()
|
|
defer d.mu.Unlock()
|
|
if d.stopped {
|
|
return ctx, nil
|
|
}
|
|
ctx, cancel := context.WithCancel(ctx)
|
|
w := &requestWork{cancel: cancel}
|
|
if d.active == nil {
|
|
d.active = make(map[*requestWork]struct{})
|
|
}
|
|
d.active[w] = struct{}{}
|
|
return ctx, w
|
|
}
|
|
|
|
func (d *requestDrain) end(w *requestWork) {
|
|
w.cancel()
|
|
d.mu.Lock()
|
|
defer d.mu.Unlock()
|
|
delete(d.active, w)
|
|
if d.stopped && len(d.active) == 0 {
|
|
close(d.done)
|
|
}
|
|
}
|
|
|
|
func (d *requestDrain) stop() {
|
|
d.mu.Lock()
|
|
defer d.mu.Unlock()
|
|
if d.stopped {
|
|
return
|
|
}
|
|
d.stopped = true
|
|
if len(d.active) == 0 {
|
|
close(d.done)
|
|
}
|
|
}
|
|
|
|
func (d *requestDrain) cancelAndWait(timeout time.Duration) error {
|
|
d.stop()
|
|
d.mu.Lock()
|
|
for w := range d.active {
|
|
w.cancel()
|
|
}
|
|
d.mu.Unlock()
|
|
timer := time.NewTimer(timeout)
|
|
defer timer.Stop()
|
|
select {
|
|
case <-d.done:
|
|
return nil
|
|
case <-timer.C:
|
|
select {
|
|
case <-d.done:
|
|
return nil
|
|
default:
|
|
}
|
|
return fmt.Errorf("go-omni: HTTP handlers did not finish within cancellation grace %s: %w", timeout,
|
|
errors.Join(context.DeadlineExceeded, lifecycle.ErrResourcesRetained))
|
|
}
|
|
}
|
|
|
|
// watchShutdown stops admission on the early signal without canceling accepted
|
|
// work. The transport retains its graceful drain budget before Close cancels.
|
|
func (a *App) watchShutdown(ctx context.Context) func() bool {
|
|
return context.AfterFunc(ctx, a.requests.stop)
|
|
}
|