| 1 | package schedule |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "errors" |
| 6 | "testing" |
| 7 | "time" |
| 8 | ) |
| 9 | |
| 10 | // A nil gate is main's runner and converger: every host enters, and a drain |
| 11 | // has nothing to wait for. |
| 12 | func TestNilHostGateIsInert(t *testing.T) { |
| 13 | var g *HostGate |
| 14 | if !g.tryEnter("a") { |
| 15 | t.Fatal("a nil gate refused a host") |
| 16 | } |
| 17 | g.leave("a") |
| 18 | if err := g.Drain(context.Background(), "a"); err != nil { |
| 19 | t.Fatalf("a nil gate's drain = %v", err) |
| 20 | } |
| 21 | } |
| 22 | |
| 23 | // The drain waits for the host's in-flight work, refuses new work while it |
| 24 | // waits, and leaves the host open once it returns. Other hosts are never |
| 25 | // held up. |
| 26 | func TestHostGateDrainWaitsForInFlightWork(t *testing.T) { |
| 27 | g := &HostGate{} |
| 28 | if !g.tryEnter("a") { |
| 29 | t.Fatal("first entry refused") |
| 30 | } |
| 31 | done := make(chan error, 1) |
| 32 | go func() { done <- g.Drain(context.Background(), "a") }() |
| 33 | select { |
| 34 | case err := <-done: |
| 35 | t.Fatalf("the drain returned (%v) while work on the host was in flight", err) |
| 36 | case <-time.After(50 * time.Millisecond): |
| 37 | } |
| 38 | // A drain is waiting: new work on the host is refused, other hosts enter. |
| 39 | deadline := time.Now().Add(time.Second) |
| 40 | for g.tryEnter("a") { |
| 41 | g.leave("a") |
| 42 | if time.Now().After(deadline) { |
| 43 | t.Fatal("work entered the host while a drain was waiting") |
| 44 | } |
| 45 | time.Sleep(time.Millisecond) |
| 46 | } |
| 47 | if !g.tryEnter("b") { |
| 48 | t.Fatal("a drain on one host refused another host") |
| 49 | } |
| 50 | g.leave("b") |
| 51 | g.leave("a") |
| 52 | if err := <-done; err != nil { |
| 53 | t.Fatalf("drain = %v", err) |
| 54 | } |
| 55 | if !g.tryEnter("a") { |
| 56 | t.Fatal("the host stayed closed after the drain") |
| 57 | } |
| 58 | g.leave("a") |
| 59 | } |
| 60 | |
| 61 | // A drain whose context ends returns the context's error; the host's work |
| 62 | // is unaffected. |
| 63 | func TestHostGateDrainHonoursItsContext(t *testing.T) { |
| 64 | g := &HostGate{} |
| 65 | if !g.tryEnter("a") { |
| 66 | t.Fatal("entry refused") |
| 67 | } |
| 68 | ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond) |
| 69 | defer cancel() |
| 70 | if err := g.Drain(ctx, "a"); !errors.Is(err, context.DeadlineExceeded) { |
| 71 | t.Fatalf("drain = %v, want the deadline", err) |
| 72 | } |
| 73 | g.leave("a") |
| 74 | } |
| 75 | |
amadan