forked from tailscale/tailscale
-
Notifications
You must be signed in to change notification settings - Fork 4
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
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> (cherry picked from commit 38a1cf7)
- Loading branch information
Showing
4 changed files
with
131 additions
and
95 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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() | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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) | ||
} | ||
} |