Add lock for storing content from source (#37)

* add lock for storing content from source

* move lock into client

* add failed_to_acquire_lock field to proto response object

* increase and make constant

* missed spot

Author: dleviminzi

pkg/client.go

@@ -540,6 +540,30 @@ func (c *BlobCacheClient) StoreContentFromSource(sourcePath string, sourceOffset
 	return resp.Hash, nil
 }
 
+func (c *BlobCacheClient) StoreContentFromSourceWithLock(sourcePath string, sourceOffset int64) (string, error) {
+	ctx, cancel := context.WithTimeout(c.ctx, storeContentRequestTimeout)
+	defer cancel()
+
+	client, _, err := c.getGRPCClient(ctx, &ClientRequest{
+		rt: ClientRequestTypeStorage,
+	})
+	if err != nil {
+		return "", err
+	}
+
+	resp, err := client.StoreContentFromSourceWithLock(ctx, &proto.StoreContentFromSourceRequest{SourcePath: sourcePath, SourceOffset: sourceOffset})
+	if err != nil {
+		return "", err
+	}
+
+	if resp.FailedToAcquireLock {
+		return "", ErrUnableToAcquireLock
+	}
+
+	return resp.Hash, nil
+
+}
+
 func (c *BlobCacheClient) HostsAvailable() bool {
 	return c.hostMap.Members().Cardinality() > 0
 }

pkg/errors.go

@@ -12,4 +12,5 @@ var (
 	ErrUnableToPopulateContent = errors.New("unable to populate content from original source")
 	ErrBlobFsMountFailure      = errors.New("failed to mount blobfs")
 	ErrTailscaleAuthentication = errors.New("tailscale authentication failed")
+	ErrUnableToAcquireLock     = errors.New("unable to acquire lock")
 )

pkg/metadata.go

@@ -16,6 +16,10 @@ import (
 	redis "github.com/redis/go-redis/v9"
 )
 
+const (
+	storeFromContentLockTtlS = 5
+)
+
 type BlobCacheMetadata struct {
 	rdb  *RedisClient
 	lock *RedisLock
@@ -368,17 +372,30 @@ func (m *BlobCacheMetadata) RemoveFsNodeChild(ctx context.Context, id string) er
 	return nil
 }
 
+func (m *BlobCacheMetadata) SetStoreFromContentLock(ctx context.Context, sourcePath string) error {
+	return m.lock.Acquire(ctx, MetadataKeys.MetadataStoreFromContentLock(sourcePath), RedisLockOptions{TtlS: storeFromContentLockTtlS, Retries: 0})
+}
+
+func (m *BlobCacheMetadata) RefreshStoreFromContentLock(ctx context.Context, sourcePath string) error {
+	return m.lock.Refresh(MetadataKeys.MetadataStoreFromContentLock(sourcePath), RedisLockOptions{TtlS: storeFromContentLockTtlS, Retries: 0})
+}
+
+func (m *BlobCacheMetadata) RemoveStoreFromContentLock(ctx context.Context, sourcePath string) error {
+	return m.lock.Release(MetadataKeys.MetadataStoreFromContentLock(sourcePath))
+}
+
 // Metadata key storage format
 var (
-	metadataPrefix         string = "blobcache"
-	metadataHostIndex      string = "blobcache:host_index"
-	metadataEntry          string = "blobcache:entry:%s"
-	metadataClientLock     string = "blobcache:client_lock:%s:%s"
-	metadataLocation       string = "blobcache:location:%s"
-	metadataRef            string = "blobcache:ref:%s"
-	metadataFsNode         string = "blobcache:fs:node:%s"
-	metadataFsNodeChildren string = "blobcache:fs:node:%s:children"
-	metadataHostKeepAlive  string = "blobcache:host:keepalive:%s"
+	metadataPrefix               string = "blobcache"
+	metadataHostIndex            string = "blobcache:host_index"
+	metadataEntry                string = "blobcache:entry:%s"
+	metadataClientLock           string = "blobcache:client_lock:%s:%s"
+	metadataLocation             string = "blobcache:location:%s"
+	metadataRef                  string = "blobcache:ref:%s"
+	metadataFsNode               string = "blobcache:fs:node:%s"
+	metadataFsNodeChildren       string = "blobcache:fs:node:%s:children"
+	metadataHostKeepAlive        string = "blobcache:host:keepalive:%s"
+	metadataStoreFromContentLock string = "blobcache:store_from_content_lock:%s"
 )
 
 // Metadata keys
@@ -418,6 +435,11 @@ func (k *metadataKeys) MetadataClientLock(hostname, hash string) string {
 	return fmt.Sprintf(metadataClientLock, hostname, hash)
 }
 
+func (k *metadataKeys) MetadataStoreFromContentLock(sourcePath string) string {
+	sourcePath = strings.ReplaceAll(sourcePath, "/", "_")
+	return fmt.Sprintf(metadataStoreFromContentLock, sourcePath)
+}
+
 var MetadataKeys = &metadataKeys{}
 
 type metadataKeys struct{}

pkg/redis.go

@@ -312,6 +312,20 @@ func (l *RedisLock) Release(key string) error {
 	return redislock.ErrLockNotHeld
 }
 
+func (l *RedisLock) Refresh(key string, opts RedisLockOptions) error {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	lock, ok := l.locks[key]
+	if !ok {
+		return redislock.ErrLockNotHeld
+	}
+
+	return lock.Refresh(context.TODO(), time.Duration(opts.TtlS)*time.Second, &redislock.Options{
+		RetryStrategy: redislock.LimitRetry(redislock.LinearBackoff(100*time.Millisecond), opts.Retries),
+	})
+}
+
 // Copies the result of HGetAll to a provided struct.
 // If a field cannot be parsed, we use Go's default value.
 // Struct fields must have the redis tag on them otherwise they will be ignored.

pkg/server.go

@@ -348,3 +348,38 @@ func (cs *CacheService) StoreContentFromSource(ctx context.Context, req *proto.S
 
 	return &proto.StoreContentFromSourceResponse{Ok: true, Hash: hash}, nil
 }
+
+func (cs *CacheService) StoreContentFromSourceWithLock(ctx context.Context, req *proto.StoreContentFromSourceRequest) (*proto.StoreContentFromSourceWithLockResponse, error) {
+	sourcePath := req.SourcePath
+	if err := cs.metadata.SetStoreFromContentLock(ctx, sourcePath); err != nil {
+		return &proto.StoreContentFromSourceWithLockResponse{FailedToAcquireLock: true}, nil
+	}
+
+	storeContext, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	go func() {
+		ticker := time.NewTicker(time.Second)
+		defer ticker.Stop()
+		for {
+			select {
+			case <-storeContext.Done():
+				return
+			case <-ticker.C:
+				Logger.Infof("StoreContentFromSourceWithLock[REFRESH] - [%s]", sourcePath)
+				cs.metadata.RefreshStoreFromContentLock(ctx, sourcePath)
+			}
+		}
+	}()
+
+	storeContentFromSourceResp, err := cs.StoreContentFromSource(storeContext, req)
+	if err != nil {
+		return nil, err
+	}
+
+	if err := cs.metadata.RemoveStoreFromContentLock(ctx, sourcePath); err != nil {
+		Logger.Errorf("StoreContentFromSourceWithLock[ERR] - error removing lock: %v", err)
+	}
+
+	return &proto.StoreContentFromSourceWithLockResponse{Hash: storeContentFromSourceResp.Hash}, nil
+}