mirror of https://github.com/tailscale/tailscale/
control/controlclient,util/execqueue: extract execqueue into a package
This is a useful primitive for asynchronous execution of ordered work I want to use in another change. Updates tailscale/corp#16833 Signed-off-by: James Tucker <james@tailscale.com>pull/10882/head
parent
32f01acc79
commit
38a1cf748a
@ -0,0 +1,104 @@
|
||||
// Copyright (c) Tailscale Inc & AUTHORS
|
||||
// SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
// Package execqueue implements an ordered asynchronous queue for executing functions.
|
||||
package execqueue
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
)
|
||||
|
||||
type ExecQueue struct {
|
||||
mu sync.Mutex
|
||||
closed bool
|
||||
inFlight bool // whether a goroutine is running q.run
|
||||
doneWaiter chan struct{} // non-nil if waiter is waiting, then closed
|
||||
queue []func()
|
||||
}
|
||||
|
||||
func (q *ExecQueue) Add(f func()) {
|
||||
q.mu.Lock()
|
||||
defer q.mu.Unlock()
|
||||
if q.closed {
|
||||
return
|
||||
}
|
||||
if q.inFlight {
|
||||
q.queue = append(q.queue, f)
|
||||
} else {
|
||||
q.inFlight = true
|
||||
go q.run(f)
|
||||
}
|
||||
}
|
||||
|
||||
// RunSync waits for the queue to be drained and then synchronously runs f.
|
||||
// It returns an error if the queue is closed before f is run or ctx expires.
|
||||
func (q *ExecQueue) RunSync(ctx context.Context, f func()) error {
|
||||
for {
|
||||
if err := q.Wait(ctx); err != nil {
|
||||
return err
|
||||
}
|
||||
q.mu.Lock()
|
||||
if q.inFlight {
|
||||
q.mu.Unlock()
|
||||
continue
|
||||
}
|
||||
defer q.mu.Unlock()
|
||||
if q.closed {
|
||||
return errors.New("closed")
|
||||
}
|
||||
f()
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func (q *ExecQueue) run(f func()) {
|
||||
f()
|
||||
|
||||
q.mu.Lock()
|
||||
for len(q.queue) > 0 && !q.closed {
|
||||
f := q.queue[0]
|
||||
q.queue[0] = nil
|
||||
q.queue = q.queue[1:]
|
||||
q.mu.Unlock()
|
||||
f()
|
||||
q.mu.Lock()
|
||||
}
|
||||
q.inFlight = false
|
||||
q.queue = nil
|
||||
if q.doneWaiter != nil {
|
||||
close(q.doneWaiter)
|
||||
q.doneWaiter = nil
|
||||
}
|
||||
q.mu.Unlock()
|
||||
}
|
||||
|
||||
// Shutdown asynchronously signals the queue to stop.
|
||||
func (q *ExecQueue) Shutdown() {
|
||||
q.mu.Lock()
|
||||
defer q.mu.Unlock()
|
||||
q.closed = true
|
||||
}
|
||||
|
||||
// Wait waits for the queue to be empty.
|
||||
func (q *ExecQueue) Wait(ctx context.Context) error {
|
||||
q.mu.Lock()
|
||||
waitCh := q.doneWaiter
|
||||
if q.inFlight && waitCh == nil {
|
||||
waitCh = make(chan struct{})
|
||||
q.doneWaiter = waitCh
|
||||
}
|
||||
q.mu.Unlock()
|
||||
|
||||
if waitCh == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
select {
|
||||
case <-waitCh:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
@ -0,0 +1,22 @@
|
||||
// Copyright (c) Tailscale Inc & AUTHORS
|
||||
// SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
package execqueue
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestExecQueue(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
var n atomic.Int32
|
||||
q := &ExecQueue{}
|
||||
defer q.Shutdown()
|
||||
q.Add(func() { n.Add(1) })
|
||||
q.Wait(ctx)
|
||||
if got := n.Load(); got != 1 {
|
||||
t.Errorf("n=%d; want 1", got)
|
||||
}
|
||||
}
|
Loading…
Reference in New Issue