github.com/grafana/pyroscope@v1.18.0/pkg/objstore/delayed_bucket_client.go (about)

     1  // SPDX-License-Identifier: AGPL-3.0-only
     2  
     3  package objstore
     4  
     5  import (
     6  	"context"
     7  	"io"
     8  	"math/rand"
     9  	"time"
    10  
    11  	"github.com/thanos-io/objstore"
    12  )
    13  
    14  // DelayedBucketClient wraps objstore.InstrumentedBucket and add a random delay to each API call.
    15  // This client is intended to be used only for testing purposes.
    16  type DelayedBucketClient struct {
    17  	wrapped  objstore.Bucket
    18  	minDelay time.Duration
    19  	maxDelay time.Duration
    20  }
    21  
    22  func NewDelayedBucketClient(wrapped objstore.Bucket, minDelay, maxDelay time.Duration) objstore.Bucket {
    23  	if minDelay < 0 || maxDelay < 0 || maxDelay < minDelay {
    24  		// We're fine just panicking, because we expect this client to be used only for testing purposes.
    25  		panic("invalid delay configuration")
    26  	}
    27  
    28  	return &DelayedBucketClient{
    29  		wrapped:  wrapped,
    30  		minDelay: minDelay,
    31  		maxDelay: maxDelay,
    32  	}
    33  }
    34  
    35  func (m *DelayedBucketClient) Upload(ctx context.Context, name string, r io.Reader, opts ...objstore.ObjectUploadOption) error {
    36  	m.delay()
    37  	defer m.delay()
    38  
    39  	return m.wrapped.Upload(ctx, name, r, opts...)
    40  }
    41  
    42  func (m *DelayedBucketClient) Delete(ctx context.Context, name string) error {
    43  	m.delay()
    44  	defer m.delay()
    45  
    46  	return m.wrapped.Delete(ctx, name)
    47  }
    48  
    49  func (m *DelayedBucketClient) Name() string {
    50  	return m.wrapped.Name()
    51  }
    52  
    53  func (m *DelayedBucketClient) Iter(ctx context.Context, dir string, f func(string) error, options ...objstore.IterOption) error {
    54  	m.delay()
    55  	defer m.delay()
    56  
    57  	return m.wrapped.Iter(ctx, dir, f, options...)
    58  }
    59  
    60  func (m *DelayedBucketClient) IterWithAttributes(ctx context.Context, dir string, f func(attrs objstore.IterObjectAttributes) error, options ...objstore.IterOption) error {
    61  	m.delay()
    62  	defer m.delay()
    63  	return m.wrapped.IterWithAttributes(ctx, dir, f, options...)
    64  }
    65  
    66  func (m *DelayedBucketClient) SupportedIterOptions() []objstore.IterOptionType {
    67  	return m.wrapped.SupportedIterOptions()
    68  }
    69  
    70  func (m *DelayedBucketClient) Get(ctx context.Context, name string) (io.ReadCloser, error) {
    71  	m.delay()
    72  	defer m.delay()
    73  
    74  	return m.wrapped.Get(ctx, name)
    75  }
    76  func (m *DelayedBucketClient) GetRange(ctx context.Context, name string, off, length int64) (io.ReadCloser, error) {
    77  	m.delay()
    78  	defer m.delay()
    79  
    80  	return m.wrapped.GetRange(ctx, name, off, length)
    81  }
    82  
    83  func (m *DelayedBucketClient) Exists(ctx context.Context, name string) (bool, error) {
    84  	m.delay()
    85  	defer m.delay()
    86  
    87  	return m.wrapped.Exists(ctx, name)
    88  }
    89  
    90  func (m *DelayedBucketClient) IsObjNotFoundErr(err error) bool {
    91  	return m.wrapped.IsObjNotFoundErr(err)
    92  }
    93  
    94  // IsAccessDeniedErr returns true if acces to object is denied.
    95  func (m *DelayedBucketClient) IsAccessDeniedErr(err error) bool {
    96  	return m.wrapped.IsAccessDeniedErr(err)
    97  }
    98  
    99  func (m *DelayedBucketClient) Attributes(ctx context.Context, name string) (objstore.ObjectAttributes, error) {
   100  	m.delay()
   101  	defer m.delay()
   102  
   103  	return m.wrapped.Attributes(ctx, name)
   104  }
   105  
   106  func (m *DelayedBucketClient) Close() error {
   107  	return m.wrapped.Close()
   108  }
   109  
   110  func (m *DelayedBucketClient) Provider() objstore.ObjProvider {
   111  	return m.wrapped.Provider()
   112  }
   113  
   114  func (m *DelayedBucketClient) delay() {
   115  	time.Sleep(m.minDelay + time.Duration(rand.Int63n(m.maxDelay.Nanoseconds()-m.minDelay.Nanoseconds())))
   116  }