github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/internal/taskrunner/preloadedtaskrunner_test.go (about)

     1  package taskrunner
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sync"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/authzed/spicedb/pkg/testutil"
    11  
    12  	"github.com/stretchr/testify/require"
    13  
    14  	"go.uber.org/goleak"
    15  )
    16  
    17  func TestPreloadedTaskRunnerCompletesAllTasks(t *testing.T) {
    18  	defer goleak.VerifyNone(t)
    19  
    20  	tr := NewPreloadedTaskRunner(context.Background(), 2, 5)
    21  	wg := sync.WaitGroup{}
    22  
    23  	for i := 0; i < 5; i++ {
    24  		wg.Add(1)
    25  		i := i
    26  		tr.Add(func(ctx context.Context) error {
    27  			time.Sleep(time.Duration(i*10) * time.Millisecond)
    28  			wg.Done()
    29  			return nil
    30  		})
    31  	}
    32  
    33  	tr.Start()
    34  
    35  	testutil.RequireWithin(t, func(t *testing.T) {
    36  		wg.Wait()
    37  	}, 5*time.Second)
    38  }
    39  
    40  func TestPreloadedTaskRunnerCancelsEarlyDueToError(t *testing.T) {
    41  	defer goleak.VerifyNone(t)
    42  
    43  	ctx := context.Background()
    44  	ctx, cancel := context.WithCancel(ctx)
    45  	defer cancel()
    46  
    47  	tr := NewPreloadedTaskRunner(ctx, 3, 10)
    48  	completed := sync.Map{}
    49  
    50  	for i := 0; i < 10; i++ {
    51  		i := i
    52  		tr.Add(func(ctx context.Context) error {
    53  			if i == 1 {
    54  				return fmt.Errorf("some error")
    55  			}
    56  
    57  			time.Sleep(time.Duration(i*50) * time.Millisecond)
    58  			completed.Store(i, true)
    59  			return nil
    60  		})
    61  	}
    62  
    63  	tr.Start()
    64  
    65  	time.Sleep(1 * time.Second)
    66  
    67  	count := 0
    68  	for i := 0; i < 10; i++ {
    69  		if _, ok := completed.Load(i); ok {
    70  			count++
    71  		}
    72  	}
    73  
    74  	require.GreaterOrEqual(t, count, 1)
    75  	require.Less(t, count, 9)
    76  }
    77  
    78  func TestPreloadedTaskRunnerCancelsEarlyDueToCancel(t *testing.T) {
    79  	defer goleak.VerifyNone(t)
    80  
    81  	ctx := context.Background()
    82  	ctx, cancel := context.WithCancel(ctx)
    83  	defer cancel()
    84  
    85  	tr := NewPreloadedTaskRunner(ctx, 3, 10)
    86  	completed := sync.Map{}
    87  
    88  	for i := 0; i < 10; i++ {
    89  		i := i
    90  		tr.Add(func(ctx context.Context) error {
    91  			if i == 1 {
    92  				cancel()
    93  				return nil
    94  			}
    95  
    96  			time.Sleep(time.Duration(i*50) * time.Millisecond)
    97  			completed.Store(i, true)
    98  			return nil
    99  		})
   100  	}
   101  
   102  	tr.Start()
   103  
   104  	time.Sleep(1 * time.Second)
   105  
   106  	count := 0
   107  	for i := 0; i < 10; i++ {
   108  		if _, ok := completed.Load(i); ok {
   109  			count++
   110  		}
   111  	}
   112  
   113  	require.GreaterOrEqual(t, count, 1)
   114  	require.Less(t, count, 9)
   115  }
   116  
   117  func TestPreloadedTaskRunnerReturnsError(t *testing.T) {
   118  	defer goleak.VerifyNone(t)
   119  
   120  	ctx := context.Background()
   121  	ctx, cancel := context.WithCancel(ctx)
   122  	defer cancel()
   123  
   124  	tr := NewPreloadedTaskRunner(ctx, 3, 10)
   125  	completed := sync.Map{}
   126  
   127  	for i := 0; i < 10; i++ {
   128  		i := i
   129  		tr.Add(func(ctx context.Context) error {
   130  			if i == 1 {
   131  				return fmt.Errorf("some error")
   132  			}
   133  
   134  			time.Sleep(time.Duration(i*50) * time.Millisecond)
   135  			completed.Store(i, true)
   136  			return nil
   137  		})
   138  	}
   139  
   140  	time.Sleep(1 * time.Second)
   141  
   142  	err := tr.StartAndWait()
   143  	require.ErrorContains(t, err, "some error")
   144  
   145  	count := 0
   146  	for i := 0; i < 10; i++ {
   147  		if _, ok := completed.Load(i); ok {
   148  			count++
   149  		}
   150  	}
   151  
   152  	require.GreaterOrEqual(t, count, 1)
   153  	require.Less(t, count, 9)
   154  }
   155  
   156  func TestPreloadedTaskRunnerEmpty(t *testing.T) {
   157  	defer goleak.VerifyNone(t)
   158  
   159  	ctx := context.Background()
   160  	ctx, cancel := context.WithCancel(ctx)
   161  	defer cancel()
   162  
   163  	tr := NewPreloadedTaskRunner(ctx, 3, 10)
   164  	err := tr.StartAndWait()
   165  	require.NoError(t, err)
   166  }