61 lines
1.6 KiB
Go
61 lines
1.6 KiB
Go
package task_test
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"git.yuntuo.hk/k2safe/go-omni/lifecycle"
|
|
"git.yuntuo.hk/k2safe/go-omni/task"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestRunnerPhasedShutdown(t *testing.T) {
|
|
runtime := lifecycle.NewRuntime(context.Background())
|
|
runner, err := task.New(task.Config{Workers: 1, QueueSize: 1})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var resourceClosed atomic.Bool
|
|
runtime.OnClose(func() error { resourceClosed.Store(true); return nil })
|
|
started, release, draining := make(chan struct{}), make(chan struct{}), make(chan struct{})
|
|
runtime.OnStop("jobs", runner.StopAccepting)
|
|
runtime.OnDrain("jobs", func(ctx context.Context) error { close(draining); return runner.Wait(ctx) })
|
|
future, err := runner.Submit(runtime.WorkContext(), "write", func(ctx context.Context) error {
|
|
close(started)
|
|
<-release
|
|
if resourceClosed.Load() {
|
|
return errors.New("resource already closed")
|
|
}
|
|
return ctx.Err()
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
<-started
|
|
runtime.RequestShutdown(nil)
|
|
done := make(chan error, 1)
|
|
go func() { done <- runtime.Close() }()
|
|
select {
|
|
case <-draining:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("did not drain")
|
|
}
|
|
if _, err := runner.Submit(context.Background(), "late", func(context.Context) error { return nil }); !errors.Is(err, task.ErrClosed) {
|
|
t.Fatal(err)
|
|
}
|
|
if resourceClosed.Load() {
|
|
t.Fatal("closed resource before task")
|
|
}
|
|
close(release)
|
|
if err := <-done; err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := future.Wait(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if !resourceClosed.Load() {
|
|
t.Fatal("resource leaked")
|
|
}
|
|
}
|