github.com/onflow/flow-go@v0.33.17/engine/access/state_stream/backend/subscription.go (about)

     1  package backend
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/google/uuid"
    10  	"google.golang.org/grpc/status"
    11  
    12  	"github.com/onflow/flow-go/engine/access/state_stream"
    13  )
    14  
    15  // GetDataByHeightFunc is a callback used by subscriptions to retrieve data for a given height.
    16  // Expected errors:
    17  // - storage.ErrNotFound
    18  // - execution_data.BlobNotFoundError
    19  // All other errors are considered exceptions
    20  type GetDataByHeightFunc func(ctx context.Context, height uint64) (interface{}, error)
    21  
    22  var _ state_stream.Subscription = (*SubscriptionImpl)(nil)
    23  
    24  type SubscriptionImpl struct {
    25  	id string
    26  
    27  	// ch is the channel used to pass data to the receiver
    28  	ch chan interface{}
    29  
    30  	// err is the error that caused the subscription to fail
    31  	err error
    32  
    33  	// once is used to ensure that the channel is only closed once
    34  	once sync.Once
    35  
    36  	// closed tracks whether or not the subscription has been closed
    37  	closed bool
    38  }
    39  
    40  func NewSubscription(bufferSize int) *SubscriptionImpl {
    41  	return &SubscriptionImpl{
    42  		id: uuid.New().String(),
    43  		ch: make(chan interface{}, bufferSize),
    44  	}
    45  }
    46  
    47  // ID returns the subscription ID
    48  // Note: this is not a cryptographic hash
    49  func (sub *SubscriptionImpl) ID() string {
    50  	return sub.id
    51  }
    52  
    53  // Channel returns the channel from which subscription data can be read
    54  func (sub *SubscriptionImpl) Channel() <-chan interface{} {
    55  	return sub.ch
    56  }
    57  
    58  // Err returns the error that caused the subscription to fail
    59  func (sub *SubscriptionImpl) Err() error {
    60  	return sub.err
    61  }
    62  
    63  // Fail registers an error and closes the subscription channel
    64  func (sub *SubscriptionImpl) Fail(err error) {
    65  	sub.err = err
    66  	sub.Close()
    67  }
    68  
    69  // Close is called when a subscription ends gracefully, and closes the subscription channel
    70  func (sub *SubscriptionImpl) Close() {
    71  	sub.once.Do(func() {
    72  		close(sub.ch)
    73  		sub.closed = true
    74  	})
    75  }
    76  
    77  // Send sends a value to the subscription channel or returns an error
    78  // Expected errors:
    79  // - context.DeadlineExceeded if send timed out
    80  // - context.Canceled if the client disconnected
    81  func (sub *SubscriptionImpl) Send(ctx context.Context, v interface{}, timeout time.Duration) error {
    82  	if sub.closed {
    83  		return fmt.Errorf("subscription closed")
    84  	}
    85  
    86  	waitCtx, cancel := context.WithTimeout(ctx, timeout)
    87  	defer cancel()
    88  
    89  	select {
    90  	case <-waitCtx.Done():
    91  		return waitCtx.Err()
    92  	case sub.ch <- v:
    93  		return nil
    94  	}
    95  }
    96  
    97  // NewFailedSubscription returns a new subscription that has already failed with the given error and
    98  // message. This is useful to return an error that occurred during subscription setup.
    99  func NewFailedSubscription(err error, msg string) *SubscriptionImpl {
   100  	sub := NewSubscription(0)
   101  
   102  	// if error is a grpc error, wrap it to preserve the error code
   103  	if st, ok := status.FromError(err); ok {
   104  		sub.Fail(status.Errorf(st.Code(), "%s: %s", msg, st.Message()))
   105  		return sub
   106  	}
   107  
   108  	// otherwise, return wrap the message normally
   109  	sub.Fail(fmt.Errorf("%s: %w", msg, err))
   110  	return sub
   111  }
   112  
   113  var _ state_stream.Subscription = (*HeightBasedSubscription)(nil)
   114  var _ state_stream.Streamable = (*HeightBasedSubscription)(nil)
   115  
   116  // HeightBasedSubscription is a subscription that retrieves data sequentially by block height
   117  type HeightBasedSubscription struct {
   118  	*SubscriptionImpl
   119  	nextHeight uint64
   120  	getData    GetDataByHeightFunc
   121  }
   122  
   123  func NewHeightBasedSubscription(bufferSize int, firstHeight uint64, getData GetDataByHeightFunc) *HeightBasedSubscription {
   124  	return &HeightBasedSubscription{
   125  		SubscriptionImpl: NewSubscription(bufferSize),
   126  		nextHeight:       firstHeight,
   127  		getData:          getData,
   128  	}
   129  }
   130  
   131  // Next returns the value for the next height from the subscription
   132  func (s *HeightBasedSubscription) Next(ctx context.Context) (interface{}, error) {
   133  	v, err := s.getData(ctx, s.nextHeight)
   134  	if err != nil {
   135  		return nil, fmt.Errorf("could not get data for height %d: %w", s.nextHeight, err)
   136  	}
   137  	s.nextHeight++
   138  	return v, nil
   139  }