github.com/openfga/openfga@v1.5.4-rc1/pkg/storage/storagewrappers/boundedconcurrency_test.go (about)

     1  package storagewrappers
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/oklog/ulid/v2"
     9  	openfgav1 "github.com/openfga/api/proto/openfga/v1"
    10  	"github.com/stretchr/testify/require"
    11  	"golang.org/x/sync/errgroup"
    12  
    13  	"github.com/openfga/openfga/internal/mocks"
    14  	"github.com/openfga/openfga/pkg/storage"
    15  	"github.com/openfga/openfga/pkg/storage/memory"
    16  	"github.com/openfga/openfga/pkg/tuple"
    17  )
    18  
    19  func TestBoundedConcurrencyWrapper(t *testing.T) {
    20  	store := ulid.Make().String()
    21  	slowBackend := mocks.NewMockSlowDataStorage(memory.New(), time.Second)
    22  
    23  	err := slowBackend.Write(context.Background(), store, []*openfgav1.TupleKeyWithoutCondition{}, []*openfgav1.TupleKey{
    24  		tuple.NewTupleKey("obj:1", "viewer", "user:anne"),
    25  	})
    26  	require.NoError(t, err)
    27  
    28  	// Create a limited tuple reader that allows 1 concurrent read a time.
    29  	limitedTupleReader := NewBoundedConcurrencyTupleReader(slowBackend, 1)
    30  
    31  	// Do reads from 4 goroutines - each should be run serially. Should be >4 seconds.
    32  	const numRoutine = 4
    33  
    34  	var wg errgroup.Group
    35  
    36  	start := time.Now()
    37  
    38  	wg.Go(func() error {
    39  		_, err := limitedTupleReader.ReadUserTuple(context.Background(), store, tuple.NewTupleKey("obj:1", "viewer", "user:anne"))
    40  		return err
    41  	})
    42  
    43  	wg.Go(func() error {
    44  		_, err := limitedTupleReader.ReadUsersetTuples(context.Background(), store, storage.ReadUsersetTuplesFilter{
    45  			Object:   "obj:1",
    46  			Relation: "viewer",
    47  		})
    48  		return err
    49  	})
    50  
    51  	wg.Go(func() error {
    52  		_, err := limitedTupleReader.Read(context.Background(), store, nil)
    53  		return err
    54  	})
    55  
    56  	wg.Go(func() error {
    57  		_, err := limitedTupleReader.ReadStartingWithUser(
    58  			context.Background(),
    59  			store,
    60  			storage.ReadStartingWithUserFilter{
    61  				UserFilter: []*openfgav1.ObjectRelation{
    62  					{
    63  						Object:   "obj",
    64  						Relation: "viewer",
    65  					},
    66  				}})
    67  		return err
    68  	})
    69  
    70  	err = wg.Wait()
    71  	require.NoError(t, err)
    72  
    73  	end := time.Now()
    74  
    75  	require.GreaterOrEqual(t, end.Sub(start), numRoutine*time.Second)
    76  }