github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/internal/taskrunner/taskrunner_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 TestTaskRunnerCompletesAllTasks(t *testing.T) {
    18  	defer goleak.VerifyNone(t)
    19  
    20  	tr := NewTaskRunner(context.Background(), 2)
    21  	completed := sync.Map{}
    22  
    23  	for i := 0; i < 5; i++ {
    24  		i := i
    25  		tr.Schedule(func(ctx context.Context) error {
    26  			time.Sleep(time.Duration(i*10) * time.Millisecond)
    27  			completed.Store(i, true)
    28  			return nil
    29  		})
    30  	}
    31  
    32  	testutil.RequireWithin(t, func(t *testing.T) {
    33  		err := tr.Wait()
    34  		require.NoError(t, err)
    35  	}, 5*time.Second)
    36  
    37  	for i := 0; i < 5; i++ {
    38  		vle, ok := completed.Load(i)
    39  		require.True(t, ok)
    40  		require.True(t, vle.(bool))
    41  	}
    42  }
    43  
    44  func TestTaskRunnerCancelsEarlyDueToError(t *testing.T) {
    45  	defer goleak.VerifyNone(t)
    46  
    47  	ctx := context.Background()
    48  	ctx, cancel := context.WithCancel(ctx)
    49  	defer cancel()
    50  
    51  	tr := NewTaskRunner(ctx, 3)
    52  	completed := sync.Map{}
    53  
    54  	for i := 0; i < 10; i++ {
    55  		i := i
    56  		tr.Schedule(func(ctx context.Context) error {
    57  			if i == 1 {
    58  				return fmt.Errorf("some error")
    59  			}
    60  
    61  			time.Sleep(time.Duration(i*50) * time.Millisecond)
    62  			completed.Store(i, true)
    63  			return nil
    64  		})
    65  	}
    66  
    67  	testutil.RequireWithin(t, func(t *testing.T) {
    68  		err := tr.Wait()
    69  		require.Error(t, err)
    70  		require.Contains(t, err.Error(), "some error")
    71  	}, 5*time.Second)
    72  
    73  	count := 0
    74  	for i := 0; i < 10; i++ {
    75  		if _, ok := completed.Load(i); ok {
    76  			count++
    77  		}
    78  	}
    79  
    80  	require.Less(t, count, 9)
    81  }
    82  
    83  func TestTaskRunnerCancelsEarlyDueToCancel(t *testing.T) {
    84  	defer goleak.VerifyNone(t)
    85  
    86  	ctx := context.Background()
    87  	ctx, cancel := context.WithCancel(ctx)
    88  	defer cancel()
    89  
    90  	tr := NewTaskRunner(ctx, 3)
    91  	completed := sync.Map{}
    92  
    93  	for i := 0; i < 10; i++ {
    94  		i := i
    95  		tr.Schedule(func(ctx context.Context) error {
    96  			if i == 1 {
    97  				cancel()
    98  				return nil
    99  			}
   100  
   101  			time.Sleep(time.Duration(i*50) * time.Millisecond)
   102  			completed.Store(i, true)
   103  			return nil
   104  		})
   105  	}
   106  
   107  	testutil.RequireWithin(t, func(t *testing.T) {
   108  		err := tr.Wait()
   109  		require.Error(t, err)
   110  		require.Contains(t, err.Error(), "canceled")
   111  	}, 5*time.Second)
   112  
   113  	count := 0
   114  	for i := 0; i < 10; i++ {
   115  		if _, ok := completed.Load(i); ok {
   116  			count++
   117  		}
   118  	}
   119  
   120  	require.Less(t, count, 9)
   121  }
   122  
   123  func TestTaskRunnerDoesNotBlockOnQueuing(t *testing.T) {
   124  	defer goleak.VerifyNone(t)
   125  
   126  	ctx := context.Background()
   127  	ctx, cancel := context.WithCancel(ctx)
   128  	defer cancel()
   129  
   130  	completed := false
   131  	tr := NewTaskRunner(ctx, 1)
   132  	tr.Schedule(func(ctx context.Context) error {
   133  		// Schedule another task which should not block.
   134  		tr.Schedule(func(ctx context.Context) error {
   135  			completed = true
   136  			return nil
   137  		})
   138  		return nil
   139  	})
   140  
   141  	testutil.RequireWithin(t, func(t *testing.T) {
   142  		err := tr.Wait()
   143  		require.NoError(t, err)
   144  		require.True(t, completed)
   145  	}, 5*time.Second)
   146  }