Add draft of throttled copy

cmd/desync/cache.go

@@ -14,6 +14,7 @@ type cacheOptions struct {
 	cache         string
 	ignoreIndexes []string
 	ignoreChunks  []string
+	throttleRateMillis int
 }
 
 func newCacheCommand(ctx context.Context) *cobra.Command {

copy.go

@@ -11,9 +11,11 @@ import (
 // store to populate a cache. If progress is provided, it'll be called when a
 // chunk has been processed. Used to draw a progress bar, can be nil.
 func Copy(ctx context.Context, ids []ChunkID, src Store, dst WriteStore, n int, pb ProgressBar) error {
+
 	in := make(chan ChunkID)
 	g, ctx := errgroup.WithContext(ctx)
 
+
 	// Setup and start the progressbar if any
 	pb.SetTotal(len(ids))
 	pb.Start()
@@ -22,6 +24,9 @@ func Copy(ctx context.Context, ids []ChunkID, src Store, dst WriteStore, n int,
 	// Start the workers
 	for i := 0; i < n; i++ {
 		g.Go(func() error {
+
+
+
 			for id := range in {
 				pb.Increment()
 				hasChunk, err := dst.HasChunk(id)
@@ -31,7 +36,10 @@ func Copy(ctx context.Context, ids []ChunkID, src Store, dst WriteStore, n int,
 				if hasChunk {
 					continue
 				}
+
+
 				chunk, err := src.GetChunk(id)
+
 				if err != nil {
 					return err
 				}

copy_test.go

@@ -0,0 +1,36 @@
+package desync
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestCopy(t *testing.T) {
+	src_store_dir := t.TempDir()
+	dst_store_dir := t.TempDir()
+
+	src_store, err := NewLocalStore(src_store_dir, StoreOptions{})
+	require.NoError(t, err)
+
+	dst_store, err := NewLocalStore(dst_store_dir, StoreOptions{})
+	require.NoError(t, err)
+
+	first_chunk_data := []byte("some data")
+	first_chunk := NewChunk(first_chunk_data)
+	first_chunk_id := first_chunk.ID()
+
+	src_store.StoreChunk(first_chunk)
+	has_the_stored_chunk, _ := src_store.HasChunk(first_chunk_id)
+	require.True(t, has_the_stored_chunk)
+
+	chunks := make([]ChunkID, 1)
+	chunks[0] = first_chunk_id
+
+	Copy(context.Background(), chunks, src_store, dst_store, 1, NewProgressBar(""))
+	require.NoError(t, err)
+	has_the_chunk, _ := dst_store.HasChunk(first_chunk_id)
+
+	require.True(t, has_the_chunk)
+}

go.mod

@@ -59,6 +59,7 @@ require (
 	golang.org/x/oauth2 v0.7.0 // indirect
 	golang.org/x/term v0.15.0 // indirect
 	golang.org/x/text v0.14.0 // indirect
+	golang.org/x/time v0.5.0 // indirect
 	golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 // indirect
 	google.golang.org/appengine v1.6.7 // indirect
 	google.golang.org/genproto v0.0.0-20230410155749-daa745c078e1 // indirect

go.sum

@@ -202,6 +202,8 @@ golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
 golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
+golang.org/x/time v0.5.0 h1:o7cqy6amK/52YcAKIPlM3a+Fpj35zvRj2TP+e1xFSfk=
+golang.org/x/time v0.5.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3HoIrodX9oNMXvdceNzlUR8zjMvY=

ratelimitstore.go

@@ -0,0 +1,77 @@
+package desync
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/pkg/errors"
+	"golang.org/x/time/rate"
+)
+
+type ThrottleOptions struct {
+	eventRate                float64
+	burstRate int
+	timeout time.Duration
+	immediateOrFail bool
+}
+
+type RateLimitedLocalStore struct {
+
+	wrappedStore WriteStore
+
+	limiter  *rate.Limiter
+
+	options ThrottleOptions
+
+}
+
+var RateLimitedExceeded = errors.New("Rate Limit Exceeded")
+
+func NewRateLimitedLocalStore(s WriteStore, options ThrottleOptions) *RateLimitedLocalStore {
+
+	limiter := rate.NewLimiter(rate.Limit(options.eventRate), options.burstRate)
+	return &RateLimitedLocalStore{wrappedStore: s,limiter: limiter, options: options }
+}
+
+func (s RateLimitedLocalStore) GetChunk(id ChunkID) (*Chunk, error) {
+
+	return s.wrappedStore.GetChunk(id)
+}
+
+func (s RateLimitedLocalStore) HasChunk(id ChunkID) (bool, error) {
+
+
+	return s.wrappedStore.HasChunk(id)
+}
+
+
+func (s RateLimitedLocalStore) StoreChunk(chunk *Chunk) error {
+
+	// This isn't ideal because what I'm really interested is in size over the wire.
+	_, err := chunk.Data()
+	if err != nil {
+		return err
+	}
+
+	//size := len(b)
+	ctx, cancel:= context.WithTimeout(context.Background(), s.options.timeout)
+	defer cancel()
+
+	if s.options.immediateOrFail{
+	   if !s.limiter.AllowN(time.Now(),1){
+			err = errors.New("Unable to immediately store")
+	   }
+	} else{
+		err = s.limiter.WaitN(ctx,1)
+	}
+
+	if err != nil {
+
+		fmt.Println("Rate limit context error:", err)
+		return RateLimitedExceeded
+	}
+
+	return s.wrappedStore.StoreChunk(chunk)
+
+}

ratelimitstore_test.go

@@ -0,0 +1,196 @@
+package desync
+
+import (
+	"context"
+	"crypto/rand"
+	"errors"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"golang.org/x/time/rate"
+)
+
+
+func NewTestRateLimitedLocalStore(t *testing.T, eventRate float64, burstRate int, timeout time.Duration, immediateOrFail bool) *RateLimitedLocalStore{
+
+	src_store_dir := t.TempDir()
+	src_store, err := NewLocalStore(src_store_dir, StoreOptions{})
+	require.NoError(t, err)
+
+	throttleOptions := ThrottleOptions{eventRate,burstRate,timeout,immediateOrFail}
+	store :=NewRateLimitedLocalStore(src_store, throttleOptions)
+	require.Equal(t,store.options.burstRate,burstRate )
+	return store
+
+}
+
+func random_data() ([]byte,error){
+
+    b := make([]byte, 16)
+    _, err := rand.Read(b)
+    if err != nil {
+        fmt.Println("Error: ", err)
+        return b,err
+    }
+
+    return b,nil
+}
+
+
+func storeLoop(t *testing.T, max int, chunk_ids []ChunkID, store RateLimitedLocalStore){
+
+	for i := 0; i < max; i++ {
+
+		data,err := random_data()
+		require.NoError(t,err)
+		chunk := NewChunk(data)
+		chunk_ids[i] = chunk.ID()
+		err  = store.StoreChunk(chunk)
+		require.Nil(t,err)
+	}
+
+
+}
+
+func chunkCheck(t *testing.T, max int, chunk_ids []ChunkID, store RateLimitedLocalStore) {
+	for i := 0; i < max; i++ {
+
+		has,err :=	store.HasChunk(chunk_ids[i])
+		require.Nil(t,err)
+		require.True(t,has)
+
+	}
+}
+
+func TestLimiter(t *testing.T){
+
+	// This is testing the framework. So not great, but I needed to get my head around what the relation
+	// between events and burst do
+	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*1000)
+	defer cancel()
+
+	limiter := rate.NewLimiter(rate.Limit(50), 50)
+	err := limiter.WaitN(ctx,1)
+	require.Nil(t, err)
+
+	limiter = rate.NewLimiter(rate.Limit(2), 1)
+	_ = limiter.WaitN(ctx,1)
+	require.Nil(t, err)
+	err = limiter.WaitN(ctx,1)
+	require.Nil(t, err)
+
+	limiter = rate.NewLimiter(rate.Limit(1), 1)
+	_ = limiter.WaitN(ctx,1)
+	require.Nil(t, err)
+	err = limiter.WaitN(ctx,1)
+	require.NotNil(t, err)
+
+
+
+
+}
+
+func TestCopyWithNoLimit(t *testing.T) {
+
+
+	throttledStore := NewTestRateLimitedLocalStore(t,1,1,time.Second*60, false)
+	chunk_data := []byte("different datas")
+	chunk := NewChunk(chunk_data)
+
+
+	start := time.Now()
+	// We start with 1 token in the bucket and replenish at 1 token per second
+	// This should take ~10 seconds.
+	// We test it takes 8s to guard against flakiness
+	for i := 0; i < 10; i++ {
+		err  := throttledStore.StoreChunk(chunk)
+		// This test will eventually fail when I get deadlines enabled
+		require.Nil(t,err)
+	}
+	finish := time.Now()
+
+	require.True(t, finish.Sub(start).Seconds() > 8)
+}
+
+
+func TestForAFullBucketNoWait(t *testing.T) {
+
+	chunk_count := 10
+	// Bucket has initial size chunk_count
+	throttledStore := NewTestRateLimitedLocalStore(t,1,chunk_count + 1,time.Second*60, false)
+
+	chunk_ids := make([]ChunkID, chunk_count)
+	start := time.Now()
+	// The bucket is full, we shouldn't have to wait
+	storeLoop(t,chunk_count,chunk_ids,*throttledStore)
+	finish := time.Now()
+	require.True(t, finish.Sub(start).Seconds() < 2)
+	chunkCheck(t,chunk_count,chunk_ids,*throttledStore)
+}
+
+func TestForAFastReplenishmentRateLittleWait(t *testing.T) {
+
+	chunk_count := 10
+	// Bucket only has one token, but we replenish chunk_count tokens every second
+	throttledStore := NewTestRateLimitedLocalStore(t,float64( chunk_count + 1),1,time.Second*60,false)
+
+	start := time.Now()
+
+
+	chunk_ids := make([]ChunkID, chunk_count)
+	storeLoop(t,chunk_count,chunk_ids,*throttledStore)
+
+	finish := time.Now()
+	require.True(t, finish.Sub(start).Seconds() < 2)
+	chunkCheck(t,chunk_count,chunk_ids,*throttledStore)
+
+
+}
+
+func TestTimeout(t *testing.T) {
+
+	// Bucket only has one token, and we replenish very slowly. We timeout, so second invocation will fail
+	throttledStore := NewTestRateLimitedLocalStore(t,float64(1) /100,1,time.Millisecond*1000, false)
+
+
+
+	data,err := random_data()
+	require.NoError(t,err)
+	chunk := NewChunk(data)
+	err  = throttledStore.StoreChunk(chunk)
+	require.Nil(t,err)
+	err  = throttledStore.StoreChunk(chunk)
+	require.NotNil(t,err)
+	require.True(t, errors.Is(err,RateLimitedExceeded))
+}
+
+func TestNoTimeout(t *testing.T) {
+	chunk_count := 10
+	// Bucket only has one token, replenish 1 per second. Timeout is 11 seconds.
+	throttledStore := NewTestRateLimitedLocalStore(t,1,1,time.Second*11, false)
+
+
+	chunk_ids := make([]ChunkID, chunk_count)
+	storeLoop(t,chunk_count,chunk_ids,*throttledStore)
+}
+
+func TestImmediateOrFail(t *testing.T) {
+
+	// Bucket only has one token, and we replenish very slowly. Second invocation will fail
+	throttledStore := NewTestRateLimitedLocalStore(t,float64(1) /100,1,time.Second*60, true)
+
+
+
+	data,err := random_data()
+	require.NoError(t,err)
+	chunk := NewChunk(data)
+
+	err  = throttledStore.StoreChunk(chunk)
+	require.Nil(t,err)
+
+	err  = throttledStore.StoreChunk(chunk)
+	require.NotNil(t,err)
+
+}