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

     1  package storagewrappers
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	openfgav1 "github.com/openfga/api/proto/openfga/v1"
     8  	"github.com/prometheus/client_golang/prometheus"
     9  	"github.com/prometheus/client_golang/prometheus/promauto"
    10  	"go.opentelemetry.io/otel/attribute"
    11  	"go.opentelemetry.io/otel/trace"
    12  
    13  	"github.com/openfga/openfga/internal/build"
    14  	"github.com/openfga/openfga/pkg/storage"
    15  	"github.com/openfga/openfga/pkg/telemetry"
    16  )
    17  
    18  const timeWaitingSpanAttribute = "time_waiting"
    19  
    20  var _ storage.RelationshipTupleReader = (*boundedConcurrencyTupleReader)(nil)
    21  
    22  var (
    23  	boundedReadDelayMsHistogram = promauto.NewHistogramVec(prometheus.HistogramOpts{
    24  		Namespace:                       build.ProjectName,
    25  		Name:                            "datastore_bounded_read_delay_ms",
    26  		Help:                            "Time spent waiting for Read, ReadUserTuple and ReadUsersetTuples calls to the datastore",
    27  		Buckets:                         []float64{1, 3, 5, 10, 25, 50, 100, 1000, 5000}, // Milliseconds. Upper bound is config.UpstreamTimeout.
    28  		NativeHistogramBucketFactor:     1.1,
    29  		NativeHistogramMaxBucketNumber:  100,
    30  		NativeHistogramMinResetDuration: time.Hour,
    31  	}, []string{"grpc_service", "grpc_method"})
    32  )
    33  
    34  type boundedConcurrencyTupleReader struct {
    35  	storage.RelationshipTupleReader
    36  	limiter chan struct{}
    37  }
    38  
    39  // NewBoundedConcurrencyTupleReader returns a wrapper over a datastore that makes sure that there are, at most,
    40  // "concurrency" concurrent calls to Read, ReadUserTuple and ReadUsersetTuples.
    41  // Consumers can then rest assured that one client will not hoard all the database connections available.
    42  func NewBoundedConcurrencyTupleReader(wrapped storage.RelationshipTupleReader, concurrency uint32) *boundedConcurrencyTupleReader {
    43  	return &boundedConcurrencyTupleReader{
    44  		RelationshipTupleReader: wrapped,
    45  		limiter:                 make(chan struct{}, concurrency),
    46  	}
    47  }
    48  
    49  // ReadUserTuple tries to return one tuple that matches the provided key exactly.
    50  func (b *boundedConcurrencyTupleReader) ReadUserTuple(
    51  	ctx context.Context,
    52  	store string,
    53  	tupleKey *openfgav1.TupleKey,
    54  ) (*openfgav1.Tuple, error) {
    55  	b.waitForLimiter(ctx)
    56  
    57  	defer func() {
    58  		<-b.limiter
    59  	}()
    60  
    61  	return b.RelationshipTupleReader.ReadUserTuple(ctx, store, tupleKey)
    62  }
    63  
    64  // Read the set of tuples associated with `store` and `TupleKey`, which may be nil or partially filled.
    65  func (b *boundedConcurrencyTupleReader) Read(ctx context.Context, store string, tupleKey *openfgav1.TupleKey) (storage.TupleIterator, error) {
    66  	b.waitForLimiter(ctx)
    67  
    68  	defer func() {
    69  		<-b.limiter
    70  	}()
    71  
    72  	return b.RelationshipTupleReader.Read(ctx, store, tupleKey)
    73  }
    74  
    75  // ReadUsersetTuples returns all userset tuples for a specified object and relation.
    76  func (b *boundedConcurrencyTupleReader) ReadUsersetTuples(
    77  	ctx context.Context,
    78  	store string,
    79  	filter storage.ReadUsersetTuplesFilter,
    80  ) (storage.TupleIterator, error) {
    81  	b.waitForLimiter(ctx)
    82  
    83  	defer func() {
    84  		<-b.limiter
    85  	}()
    86  
    87  	return b.RelationshipTupleReader.ReadUsersetTuples(ctx, store, filter)
    88  }
    89  
    90  // ReadStartingWithUser performs a reverse read of relationship tuples starting at one or
    91  // more user(s) or userset(s) and filtered by object type and relation.
    92  func (b *boundedConcurrencyTupleReader) ReadStartingWithUser(
    93  	ctx context.Context,
    94  	store string,
    95  	filter storage.ReadStartingWithUserFilter,
    96  ) (storage.TupleIterator, error) {
    97  	b.waitForLimiter(ctx)
    98  
    99  	defer func() {
   100  		<-b.limiter
   101  	}()
   102  
   103  	return b.RelationshipTupleReader.ReadStartingWithUser(ctx, store, filter)
   104  }
   105  
   106  func (b *boundedConcurrencyTupleReader) waitForLimiter(ctx context.Context) {
   107  	start := time.Now()
   108  
   109  	b.limiter <- struct{}{}
   110  
   111  	end := time.Now()
   112  	timeWaiting := end.Sub(start).Milliseconds()
   113  
   114  	rpcInfo := telemetry.RPCInfoFromContext(ctx)
   115  	boundedReadDelayMsHistogram.WithLabelValues(
   116  		rpcInfo.Service,
   117  		rpcInfo.Method,
   118  	).Observe(float64(timeWaiting))
   119  
   120  	span := trace.SpanFromContext(ctx)
   121  	span.SetAttributes(attribute.Int64(timeWaitingSpanAttribute, timeWaiting))
   122  }