vendor tsgo

tools/tsgo/internal/core/workgroup.go

@@ -0,0 +1,123 @@
+package core
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+
+	"golang.org/x/sync/errgroup"
+)
+
+type WorkGroup interface {
+	// Queue queues a function to run. It may be invoked immediately, or deferred until RunAndWait.
+	// It is not safe to call Queue after RunAndWait has returned.
+	Queue(fn func())
+
+	// RunAndWait runs all queued functions, blocking until they have all completed.
+	RunAndWait()
+}
+
+func NewWorkGroup(singleThreaded bool) WorkGroup {
+	if singleThreaded {
+		return &singleThreadedWorkGroup{}
+	}
+	return &parallelWorkGroup{}
+}
+
+type parallelWorkGroup struct {
+	done atomic.Bool
+	wg   sync.WaitGroup
+}
+
+var _ WorkGroup = (*parallelWorkGroup)(nil)
+
+func (w *parallelWorkGroup) Queue(fn func()) {
+	if w.done.Load() {
+		panic("Queue called after RunAndWait returned")
+	}
+
+	w.wg.Go(func() {
+		fn()
+	})
+}
+
+func (w *parallelWorkGroup) RunAndWait() {
+	defer w.done.Store(true)
+	w.wg.Wait()
+}
+
+type singleThreadedWorkGroup struct {
+	done  atomic.Bool
+	fnsMu sync.Mutex
+	fns   []func()
+}
+
+var _ WorkGroup = (*singleThreadedWorkGroup)(nil)
+
+func (w *singleThreadedWorkGroup) Queue(fn func()) {
+	if w.done.Load() {
+		panic("Queue called after RunAndWait returned")
+	}
+
+	w.fnsMu.Lock()
+	defer w.fnsMu.Unlock()
+	w.fns = append(w.fns, fn)
+}
+
+func (w *singleThreadedWorkGroup) RunAndWait() {
+	defer w.done.Store(true)
+	for {
+		fn := w.pop()
+		if fn == nil {
+			return
+		}
+		fn()
+	}
+}
+
+func (w *singleThreadedWorkGroup) pop() func() {
+	w.fnsMu.Lock()
+	defer w.fnsMu.Unlock()
+	if len(w.fns) == 0 {
+		return nil
+	}
+	end := len(w.fns) - 1
+	fn := w.fns[end]
+	w.fns[end] = nil // Allow GC
+	w.fns = w.fns[:end]
+	return fn
+}
+
+// ThrottleGroup is like errgroup.Group but with global concurrency limiting via a semaphore.
+type ThrottleGroup struct {
+	semaphore chan struct{}
+	group     *errgroup.Group
+}
+
+// NewThrottleGroup creates a new ThrottleGroup with the given context and semaphore for concurrency limiting.
+func NewThrottleGroup(ctx context.Context, semaphore chan struct{}) *ThrottleGroup {
+	g, _ := errgroup.WithContext(ctx)
+	return &ThrottleGroup{
+		semaphore: semaphore,
+		group:     g,
+	}
+}
+
+// Go runs the given function in a new goroutine, but first acquires a slot from the semaphore.
+// The semaphore slot is released when the function completes.
+func (tg *ThrottleGroup) Go(fn func() error) {
+	tg.group.Go(func() error {
+		// Acquire semaphore slot - this will block until a slot is available
+		tg.semaphore <- struct{}{}
+		defer func() {
+			// Release semaphore slot when done
+			<-tg.semaphore
+		}()
+		return fn()
+	})
+}
+
+// Wait waits for all goroutines to complete and returns the first error encountered, if any.
+func (tg *ThrottleGroup) Wait() error {
+	return tg.group.Wait()
+}