niki/vendor/golang.org/x/net/http2/testsync.go

// Copyright 2024 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package http2

import (
	"context"
	"sync"
	"time"
)

// testSyncHooks coordinates goroutines in tests.
//
// For example, a call to ClientConn.RoundTrip involves several goroutines, including:
//   - the goroutine running RoundTrip;
//   - the clientStream.doRequest goroutine, which writes the request; and
//   - the clientStream.readLoop goroutine, which reads the response.
//
// Using testSyncHooks, a test can start a RoundTrip and identify when all these goroutines
// are blocked waiting for some condition such as reading the Request.Body or waiting for
// flow control to become available.
//
// The testSyncHooks also manage timers and synthetic time in tests.
// This permits us to, for example, start a request and cause it to time out waiting for
// response headers without resorting to time.Sleep calls.
type testSyncHooks struct {
	// active/inactive act as a mutex and condition variable.
	//
	//  - neither chan contains a value: testSyncHooks is locked.
	//  - active contains a value: unlocked, and at least one goroutine is not blocked
	//  - inactive contains a value: unlocked, and all goroutines are blocked
	active   chan struct{}
	inactive chan struct{}

	// goroutine counts
	total    int                     // total goroutines
	condwait map[*sync.Cond]int      // blocked in sync.Cond.Wait
	blocked  []*testBlockedGoroutine // otherwise blocked

	// fake time
	now    time.Time
	timers []*fakeTimer

	// Transport testing: Report various events.
	newclientconn func(*ClientConn)
	newstream     func(*clientStream)
}

// testBlockedGoroutine is a blocked goroutine.
type testBlockedGoroutine struct {
	f  func() bool   // blocked until f returns true
	ch chan struct{} // closed when unblocked
}

func newTestSyncHooks() *testSyncHooks {
	h := &testSyncHooks{
		active:   make(chan struct{}, 1),
		inactive: make(chan struct{}, 1),
		condwait: map[*sync.Cond]int{},
	}
	h.inactive <- struct{}{}
	h.now = time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC)
	return h
}

// lock acquires the testSyncHooks mutex.
func (h *testSyncHooks) lock() {
	select {
	case <-h.active:
	case <-h.inactive:
	}
}

// waitInactive waits for all goroutines to become inactive.
func (h *testSyncHooks) waitInactive() {
	for {
		<-h.inactive
		if !h.unlock() {
			break
		}
	}
}

// unlock releases the testSyncHooks mutex.
// It reports whether any goroutines are active.
func (h *testSyncHooks) unlock() (active bool) {
	// Look for a blocked goroutine which can be unblocked.
	blocked := h.blocked[:0]
	unblocked := false
	for _, b := range h.blocked {
		if !unblocked && b.f() {
			unblocked = true
			close(b.ch)
		} else {
			blocked = append(blocked, b)
		}
	}
	h.blocked = blocked

	// Count goroutines blocked on condition variables.
	condwait := 0
	for _, count := range h.condwait {
		condwait += count
	}

	if h.total > condwait+len(blocked) {
		h.active <- struct{}{}
		return true
	} else {
		h.inactive <- struct{}{}
		return false
	}
}

// goRun starts a new goroutine.
func (h *testSyncHooks) goRun(f func()) {
	h.lock()
	h.total++
	h.unlock()
	go func() {
		defer func() {
			h.lock()
			h.total--
			h.unlock()
		}()
		f()
	}()
}

// blockUntil indicates that a goroutine is blocked waiting for some condition to become true.
// It waits until f returns true before proceeding.
//
// Example usage:
//
//	h.blockUntil(func() bool {
//		// Is the context done yet?
//		select {
//		case <-ctx.Done():
//		default:
//			return false
//		}
//		return true
//	})
//	// Wait for the context to become done.
//	<-ctx.Done()
//
// The function f passed to blockUntil must be non-blocking and idempotent.
func (h *testSyncHooks) blockUntil(f func() bool) {
	if f() {
		return
	}
	ch := make(chan struct{})
	h.lock()
	h.blocked = append(h.blocked, &testBlockedGoroutine{
		f:  f,
		ch: ch,
	})
	h.unlock()
	<-ch
}

// broadcast is sync.Cond.Broadcast.
func (h *testSyncHooks) condBroadcast(cond *sync.Cond) {
	h.lock()
	delete(h.condwait, cond)
	h.unlock()
	cond.Broadcast()
}

// broadcast is sync.Cond.Wait.
func (h *testSyncHooks) condWait(cond *sync.Cond) {
	h.lock()
	h.condwait[cond]++
	h.unlock()
}

// newTimer creates a new fake timer.
func (h *testSyncHooks) newTimer(d time.Duration) timer {
	h.lock()
	defer h.unlock()
	t := &fakeTimer{
		hooks: h,
		when:  h.now.Add(d),
		c:     make(chan time.Time),
	}
	h.timers = append(h.timers, t)
	return t
}

// afterFunc creates a new fake AfterFunc timer.
func (h *testSyncHooks) afterFunc(d time.Duration, f func()) timer {
	h.lock()
	defer h.unlock()
	t := &fakeTimer{
		hooks: h,
		when:  h.now.Add(d),
		f:     f,
	}
	h.timers = append(h.timers, t)
	return t
}

func (h *testSyncHooks) contextWithTimeout(ctx context.Context, d time.Duration) (context.Context, context.CancelFunc) {
	ctx, cancel := context.WithCancel(ctx)
	t := h.afterFunc(d, cancel)
	return ctx, func() {
		t.Stop()
		cancel()
	}
}

func (h *testSyncHooks) timeUntilEvent() time.Duration {
	h.lock()
	defer h.unlock()
	var next time.Time
	for _, t := range h.timers {
		if next.IsZero() || t.when.Before(next) {
			next = t.when
		}
	}
	if d := next.Sub(h.now); d > 0 {
		return d
	}
	return 0
}

// advance advances time and causes synthetic timers to fire.
func (h *testSyncHooks) advance(d time.Duration) {
	h.lock()
	defer h.unlock()
	h.now = h.now.Add(d)
	timers := h.timers[:0]
	for _, t := range h.timers {
		t := t // remove after go.mod depends on go1.22
		t.mu.Lock()
		switch {
		case t.when.After(h.now):
			timers = append(timers, t)
		case t.when.IsZero():
			// stopped timer
		default:
			t.when = time.Time{}
			if t.c != nil {
				close(t.c)
			}
			if t.f != nil {
				h.total++
				go func() {
					defer func() {
						h.lock()
						h.total--
						h.unlock()
					}()
					t.f()
				}()
			}
		}
		t.mu.Unlock()
	}
	h.timers = timers
}

// A timer wraps a time.Timer, or a synthetic equivalent in tests.
// Unlike time.Timer, timer is single-use: The timer channel is closed when the timer expires.
type timer interface {
	C() <-chan time.Time
	Stop() bool
	Reset(d time.Duration) bool
}

// timeTimer implements timer using real time.
type timeTimer struct {
	t *time.Timer
	c chan time.Time
}

// newTimeTimer creates a new timer using real time.
func newTimeTimer(d time.Duration) timer {
	ch := make(chan time.Time)
	t := time.AfterFunc(d, func() {
		close(ch)
	})
	return &timeTimer{t, ch}
}

// newTimeAfterFunc creates an AfterFunc timer using real time.
func newTimeAfterFunc(d time.Duration, f func()) timer {
	return &timeTimer{
		t: time.AfterFunc(d, f),
	}
}

func (t timeTimer) C() <-chan time.Time        { return t.c }
func (t timeTimer) Stop() bool                 { return t.t.Stop() }
func (t timeTimer) Reset(d time.Duration) bool { return t.t.Reset(d) }

// fakeTimer implements timer using fake time.
type fakeTimer struct {
	hooks *testSyncHooks

	mu   sync.Mutex
	when time.Time      // when the timer will fire
	c    chan time.Time // closed when the timer fires; mutually exclusive with f
	f    func()         // called when the timer fires; mutually exclusive with c
}

func (t *fakeTimer) C() <-chan time.Time { return t.c }

func (t *fakeTimer) Stop() bool {
	t.mu.Lock()
	defer t.mu.Unlock()
	stopped := t.when.IsZero()
	t.when = time.Time{}
	return stopped
}

func (t *fakeTimer) Reset(d time.Duration) bool {
	if t.c != nil || t.f == nil {
		panic("fakeTimer only supports Reset on AfterFunc timers")
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	t.hooks.lock()
	defer t.hooks.unlock()
	active := !t.when.IsZero()
	t.when = t.hooks.now.Add(d)
	if !active {
		t.hooks.timers = append(t.hooks.timers, t)
	}
	return active
}