mirror of
https://github.com/lightningnetwork/lnd.git
synced 2025-02-23 06:35:07 +01:00
Simplifies context cancellation handling by using context.AfterFunc instead of a goroutine to wait for context cancellation. This approach avoids the overhead of a goroutine during the waiting period. For ctxQuitUnsafe, since g.quit is closed only in the Quit method (which also cancels all associated contexts), waiting on context cancellation ensures the same behavior without unnecessary dependency on g.quit. Added a test to ensure that the Create method does not launch any goroutines.
477 lines
11 KiB
Go
477 lines
11 KiB
Go
package fn
|
|
|
|
import (
|
|
"context"
|
|
"runtime"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
// TestContextGuard tests the behaviour of the ContextGuard.
|
|
func TestContextGuard(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
defer GuardTest(t)()
|
|
|
|
// Test that the derived context is cancelled when the passed context is
|
|
// cancelled.
|
|
t.Run("Parent context is cancelled", func(t *testing.T) {
|
|
t.Parallel()
|
|
var (
|
|
ctx, cancel = context.WithCancel(context.Background())
|
|
g = NewContextGuard()
|
|
)
|
|
|
|
ctxc, _ := g.Create(ctx)
|
|
|
|
// Cancel the parent context.
|
|
cancel()
|
|
// Assert that the derived context is cancelled.
|
|
select {
|
|
case <-ctxc.Done():
|
|
default:
|
|
t.Errorf("The derived context should be cancelled at " +
|
|
"this point")
|
|
}
|
|
})
|
|
|
|
// Test that the derived context is cancelled when the returned cancel
|
|
// function is called.
|
|
t.Run("Derived context is cancelled", func(t *testing.T) {
|
|
t.Parallel()
|
|
var (
|
|
ctx = context.Background()
|
|
g = NewContextGuard()
|
|
)
|
|
|
|
ctxc, cancel := g.Create(ctx)
|
|
|
|
// Cancel the context.
|
|
cancel()
|
|
|
|
// Assert that the derived context is cancelled.
|
|
select {
|
|
case <-ctxc.Done():
|
|
default:
|
|
t.Errorf("The derived context should be cancelled at " +
|
|
"this point")
|
|
}
|
|
})
|
|
|
|
// Test that the derived context is cancelled when the quit channel is
|
|
// closed.
|
|
t.Run("Quit channel is closed", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
var (
|
|
ctx = context.Background()
|
|
g = NewContextGuard()
|
|
)
|
|
ctxc, _ := g.Create(ctx)
|
|
|
|
// Close the quit channel.
|
|
g.Quit()
|
|
|
|
// Assert that the derived context is cancelled.
|
|
select {
|
|
case <-ctxc.Done():
|
|
default:
|
|
t.Errorf("The derived context should be cancelled at " +
|
|
"this point")
|
|
}
|
|
})
|
|
|
|
t.Run("Parent context is already closed", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
var (
|
|
ctx, cancel = context.WithCancel(context.Background())
|
|
g = NewContextGuard()
|
|
)
|
|
cancel()
|
|
|
|
ctxc, _ := g.Create(ctx)
|
|
|
|
// Assert that the derived context is already cancelled.
|
|
select {
|
|
case <-ctxc.Done():
|
|
default:
|
|
t.Errorf("The derived context should be cancelled at " +
|
|
"this point")
|
|
}
|
|
})
|
|
|
|
t.Run("Quit channel is already closed", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
var (
|
|
ctx = context.Background()
|
|
g = NewContextGuard()
|
|
)
|
|
|
|
g.Quit()
|
|
|
|
ctxc, _ := g.Create(ctx)
|
|
|
|
// Assert that the derived context is already cancelled.
|
|
select {
|
|
case <-ctxc.Done():
|
|
default:
|
|
t.Errorf("The derived context should be cancelled at " +
|
|
"this point")
|
|
}
|
|
})
|
|
|
|
t.Run("Child context should be cancelled synchronously with "+
|
|
"parent", func(t *testing.T) {
|
|
|
|
t.Parallel()
|
|
|
|
var (
|
|
ctx, cancel = context.WithCancel(context.Background())
|
|
g = NewContextGuard()
|
|
task = make(chan struct{})
|
|
done = make(chan struct{})
|
|
)
|
|
// Derive a child context.
|
|
ctxc, _ := g.Create(ctx)
|
|
|
|
// Spin off a routine that exists cleanly if the child context
|
|
// is cancelled but fails if the task is performed.
|
|
go func() {
|
|
defer close(done)
|
|
select {
|
|
case <-ctxc.Done():
|
|
case <-task:
|
|
t.Fatalf("should not get here")
|
|
}
|
|
}()
|
|
|
|
// Give the goroutine above a chance to spin up so that it's
|
|
// waiting on the select.
|
|
time.Sleep(time.Millisecond * 200)
|
|
|
|
// First cancel the parent context. Then immediately execute the
|
|
// task.
|
|
cancel()
|
|
close(task)
|
|
|
|
// Wait for the goroutine to exit.
|
|
select {
|
|
case <-done:
|
|
case <-time.After(time.Second):
|
|
t.Fatalf("timeout")
|
|
}
|
|
})
|
|
|
|
t.Run("Child context should be cancelled synchronously with the "+
|
|
"close of the quit channel", func(t *testing.T) {
|
|
|
|
t.Parallel()
|
|
|
|
var (
|
|
ctx = context.Background()
|
|
g = NewContextGuard()
|
|
task = make(chan struct{})
|
|
done = make(chan struct{})
|
|
)
|
|
|
|
// Derive a child context.
|
|
ctxc, _ := g.Create(ctx)
|
|
|
|
// Spin off a routine that exists cleanly if the child context
|
|
// is cancelled but fails if the task is performed.
|
|
go func() {
|
|
defer close(done)
|
|
select {
|
|
case <-ctxc.Done():
|
|
case <-task:
|
|
t.Fatalf("should not get here")
|
|
}
|
|
}()
|
|
|
|
// Give the goroutine above a chance to spin up so that it's
|
|
// waiting on the select.
|
|
time.Sleep(time.Millisecond * 200)
|
|
|
|
// First cancel the parent context. Then immediately execute the
|
|
// task.
|
|
g.Quit()
|
|
|
|
// Execute the task.
|
|
close(task)
|
|
|
|
// Wait for the goroutine to exit.
|
|
select {
|
|
case <-done:
|
|
case <-time.After(time.Second):
|
|
t.Fatalf("timeout")
|
|
}
|
|
})
|
|
|
|
// Test that if we add the BlockingCGOpt option, then the context will
|
|
// not be cancelled when the quit channel is closed but will be when the
|
|
// cancel function is called.
|
|
t.Run("Blocking context no timeout", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
var (
|
|
ctx = context.Background()
|
|
g = NewContextGuard()
|
|
task = make(chan struct{})
|
|
done = make(chan struct{})
|
|
)
|
|
|
|
// Derive a blocking child context.
|
|
ctxc, cancel := g.Create(ctx, WithBlockingCG())
|
|
|
|
// Spin of a routine that will exit cleanly if the context is
|
|
// cancelled but will fail if the task is performed.
|
|
go func() {
|
|
defer func() {
|
|
done <- struct{}{}
|
|
}()
|
|
|
|
select {
|
|
case <-ctxc.Done():
|
|
case <-task:
|
|
t.Fatalf("Expected context to be cancelled")
|
|
}
|
|
}()
|
|
|
|
// Give the goroutine above a chance to spin up so that it's
|
|
// waiting on the select.
|
|
time.Sleep(time.Millisecond * 200)
|
|
|
|
// Cancel the context.
|
|
cancel()
|
|
|
|
// Attempt to perform the task.
|
|
select {
|
|
case task <- struct{}{}:
|
|
t.Fatalf("Expected task to not be performed")
|
|
default:
|
|
}
|
|
|
|
// Assert that the task goroutine has now completed.
|
|
select {
|
|
case <-done:
|
|
case <-time.After(time.Second):
|
|
t.Fatalf("timeout")
|
|
}
|
|
|
|
// Derive a new blocking child context.
|
|
ctxc, cancel = g.Create(ctx, WithBlockingCG())
|
|
|
|
// Repeat the task but this time, we will call Quit first, but
|
|
// since this is a blocking context, the context should not be
|
|
// cancelled and the task _should_ be performed.
|
|
go func() {
|
|
defer func() {
|
|
done <- struct{}{}
|
|
}()
|
|
|
|
select {
|
|
case <-ctxc.Done():
|
|
t.Fatalf("Expected task to be performed")
|
|
case <-task:
|
|
}
|
|
}()
|
|
|
|
// Give the goroutine above a chance to spin up so that it's
|
|
// waiting on the select.
|
|
time.Sleep(time.Millisecond * 200)
|
|
|
|
// Close the quit channel. This should NOT cause the context
|
|
// to be cancelled.
|
|
g.Quit()
|
|
|
|
// Now, perform the task.
|
|
select {
|
|
case task <- struct{}{}:
|
|
case <-time.After(time.Second):
|
|
t.Fatalf("timeout")
|
|
}
|
|
|
|
// Assert that the task goroutine has now completed.
|
|
select {
|
|
case <-done:
|
|
case <-time.After(time.Second):
|
|
t.Fatalf("timeout")
|
|
}
|
|
|
|
// Cancel the context.
|
|
cancel()
|
|
|
|
// Make sure wg's counter gets to 0 eventually.
|
|
g.WgWait()
|
|
})
|
|
|
|
// Test that if we add the CustomTimeoutCGOpt option, then the context
|
|
// will be not be cancelled when the quit channel is closed but will be
|
|
// if either the context is cancelled or the timeout is reached.
|
|
t.Run("Blocking context with timeout", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
var (
|
|
ctx = context.Background()
|
|
g = NewContextGuard()
|
|
task = make(chan struct{})
|
|
done = make(chan struct{})
|
|
timeout = time.Millisecond * 500
|
|
)
|
|
|
|
// Derive a blocking child context.
|
|
ctxc, cancel := g.Create(
|
|
ctx, WithBlockingCG(), WithCustomTimeoutCG(timeout),
|
|
)
|
|
|
|
// Spin of a routine that will exit cleanly if the context is
|
|
// cancelled but will fail if the task is performed.
|
|
go func() {
|
|
defer func() {
|
|
done <- struct{}{}
|
|
}()
|
|
|
|
select {
|
|
case <-ctxc.Done():
|
|
case <-task:
|
|
t.Fatalf("Expected context to be cancelled")
|
|
}
|
|
}()
|
|
|
|
// Give the goroutine above a chance to spin up so that it's
|
|
// waiting on the select.
|
|
time.Sleep(time.Millisecond * 200)
|
|
|
|
// Cancel the context.
|
|
cancel()
|
|
|
|
// Attempt to perform the task.
|
|
select {
|
|
case task <- struct{}{}:
|
|
t.Fatalf("Expected task to not be performed")
|
|
default:
|
|
}
|
|
|
|
// Assert that the task goroutine has now completed.
|
|
select {
|
|
case <-done:
|
|
case <-time.After(time.Second):
|
|
t.Fatalf("timeout")
|
|
}
|
|
|
|
// Derive a new blocking child context with a timeout.
|
|
ctxc, cancel = g.Create(
|
|
ctx, WithBlockingCG(), WithCustomTimeoutCG(timeout),
|
|
)
|
|
|
|
// Repeat the task but this time, but this time, we will assert
|
|
// that the context is cancelled if the timeout is reached.
|
|
// We will again fail if the task is performed.
|
|
go func() {
|
|
defer func() {
|
|
done <- struct{}{}
|
|
}()
|
|
|
|
select {
|
|
case <-ctxc.Done():
|
|
case <-task:
|
|
t.Fatalf("Expected context to be cancelled")
|
|
}
|
|
}()
|
|
|
|
// Wait for the timeout to be reached.
|
|
time.Sleep(timeout + time.Millisecond*100)
|
|
|
|
// Attempt to perform the task.
|
|
select {
|
|
case task <- struct{}{}:
|
|
t.Fatalf("Expected task to not be performed")
|
|
default:
|
|
}
|
|
|
|
// Assert that the task goroutine has now completed.
|
|
select {
|
|
case <-done:
|
|
case <-time.After(time.Second):
|
|
t.Fatalf("timeout")
|
|
}
|
|
|
|
// Finally, repeat the task but this time show that calling
|
|
// Quit does not cancel the context and that the task still gets
|
|
// performed if it takes place before the context is timed out.
|
|
ctxc, cancel = g.Create(
|
|
ctx, WithBlockingCG(), WithCustomTimeoutCG(timeout),
|
|
)
|
|
|
|
go func() {
|
|
defer func() {
|
|
done <- struct{}{}
|
|
}()
|
|
|
|
select {
|
|
case <-ctxc.Done():
|
|
t.Fatalf("Expected the task to be performed")
|
|
case <-task:
|
|
}
|
|
}()
|
|
|
|
// Give the goroutine above a chance to spin up so that it's
|
|
// waiting on the select.
|
|
time.Sleep(time.Millisecond * 200)
|
|
|
|
// Close the quit channel. This should NOT cause the context
|
|
// to be cancelled.
|
|
g.Quit()
|
|
|
|
// Now, perform the task.
|
|
select {
|
|
case task <- struct{}{}:
|
|
case <-time.After(time.Second):
|
|
t.Fatalf("timeout")
|
|
}
|
|
|
|
// Assert that the task goroutine has now completed.
|
|
select {
|
|
case <-done:
|
|
case <-time.After(time.Second):
|
|
t.Fatalf("timeout")
|
|
}
|
|
})
|
|
}
|
|
|
|
// TestContextGuardCountGoroutines makes sure that ContextGuard doesn't create
|
|
// any goroutines while waiting for contexts.
|
|
func TestContextGuardCountGoroutines(t *testing.T) {
|
|
// NOTE: t.Parallel() is not called in this test because it relies on an
|
|
// accurate count of active goroutines. Running other tests in parallel
|
|
// would introduce additional goroutines, leading to unreliable results.
|
|
|
|
g := NewContextGuard()
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
|
|
// Count goroutines before contexts are created.
|
|
count1 := runtime.NumGoroutine()
|
|
|
|
// Create 1000 contexts of each type.
|
|
for i := 0; i < 1000; i++ {
|
|
_, _ = g.Create(ctx)
|
|
_, _ = g.Create(ctx, WithBlockingCG())
|
|
_, _ = g.Create(ctx, WithTimeoutCG())
|
|
_, _ = g.Create(ctx, WithBlockingCG(), WithTimeoutCG())
|
|
}
|
|
|
|
// Make sure no new goroutine was launched.
|
|
count2 := runtime.NumGoroutine()
|
|
require.LessOrEqual(t, count2, count1)
|
|
|
|
// Cancel root context.
|
|
cancel()
|
|
|
|
// Make sure wg's counter gets to 0 eventually.
|
|
g.WgWait()
|
|
}
|