361 lines
13 KiB
C#
361 lines
13 KiB
C#
// Copyright Epic Games, Inc. All Rights Reserved.
|
|
|
|
using System;
|
|
using System.Collections.Generic;
|
|
using System.Linq;
|
|
using System.Runtime.CompilerServices;
|
|
using System.Threading;
|
|
using System.Threading.Tasks;
|
|
using EpicGames.Horde.Storage;
|
|
using Microsoft.Extensions.Options;
|
|
using MongoDB.Bson.Serialization.Attributes;
|
|
using MongoDB.Bson.Serialization.Options;
|
|
using MongoDB.Driver;
|
|
|
|
namespace Jupiter.Implementation.Blob;
|
|
|
|
public class MongoBlobIndex : MongoStore, IBlobIndex
|
|
{
|
|
private readonly IOptionsMonitor<JupiterSettings> _jupiterSettings;
|
|
|
|
public MongoBlobIndex(IOptionsMonitor<JupiterSettings> jupiterSettings, IOptionsMonitor<MongoSettings> settings) : base(settings)
|
|
{
|
|
_jupiterSettings = jupiterSettings;
|
|
|
|
CreateCollectionIfNotExistsAsync<MongoBlobIndexModelV0>().Wait();
|
|
|
|
IndexKeysDefinitionBuilder<MongoBlobIndexModelV0> indexKeysDefinitionBuilder = Builders<MongoBlobIndexModelV0>.IndexKeys;
|
|
CreateIndexModel<MongoBlobIndexModelV0> indexModel = new CreateIndexModel<MongoBlobIndexModelV0>(
|
|
indexKeysDefinitionBuilder.Combine(
|
|
indexKeysDefinitionBuilder.Ascending(m => m.Ns),
|
|
indexKeysDefinitionBuilder.Ascending(m => m.BlobId)
|
|
)
|
|
, new CreateIndexOptions()
|
|
{
|
|
Name = "CompoundIndex"
|
|
});
|
|
|
|
AddIndexFor<MongoBlobIndexModelV0>().CreateMany(new[] {
|
|
indexModel,
|
|
});
|
|
}
|
|
|
|
private async Task<MongoBlobIndexModelV0?> GetBlobInfoAsync(NamespaceId ns, BlobId id, CancellationToken cancellationToken)
|
|
{
|
|
IMongoCollection<MongoBlobIndexModelV0> collection = GetCollection<MongoBlobIndexModelV0>();
|
|
IAsyncCursor<MongoBlobIndexModelV0>? cursor = await collection.FindAsync(m => m.Ns == ns.ToString() && m.BlobId == id.ToString(), cancellationToken: cancellationToken);
|
|
MongoBlobIndexModelV0? model = await cursor.FirstOrDefaultAsync(cancellationToken);
|
|
if (model == null)
|
|
{
|
|
return null;
|
|
}
|
|
|
|
return model;
|
|
}
|
|
|
|
public async Task AddBlobToIndexAsync(NamespaceId ns, BlobId id, string? region = null, CancellationToken cancellationToken = default)
|
|
{
|
|
region ??= _jupiterSettings.CurrentValue.CurrentSite;
|
|
IMongoCollection<MongoBlobIndexModelV0> collection = GetCollection<MongoBlobIndexModelV0>();
|
|
MongoBlobIndexModelV0 model = new MongoBlobIndexModelV0(ns, id);
|
|
model.Regions.Add(region);
|
|
|
|
FilterDefinition<MongoBlobIndexModelV0> filter = Builders<MongoBlobIndexModelV0>.Filter.Where(m => m.Ns == ns.ToString() && m.BlobId == id.ToString());
|
|
await collection.FindOneAndReplaceAsync(filter, model, new FindOneAndReplaceOptions<MongoBlobIndexModelV0, MongoBlobIndexModelV0>
|
|
{
|
|
IsUpsert = true
|
|
}, cancellationToken);
|
|
}
|
|
|
|
public async Task RemoveBlobFromRegionAsync(NamespaceId ns, BlobId id, string? region = null, CancellationToken cancellationToken = default)
|
|
{
|
|
region ??= _jupiterSettings.CurrentValue.CurrentSite;
|
|
|
|
MongoBlobIndexModelV0? model = await GetBlobInfoAsync(ns, id, cancellationToken);
|
|
if (model == null)
|
|
{
|
|
throw new BlobNotFoundException(ns, id);
|
|
}
|
|
IMongoCollection<MongoBlobIndexModelV0> collection = GetCollection<MongoBlobIndexModelV0>();
|
|
model.Regions.Remove(region);
|
|
FilterDefinition<MongoBlobIndexModelV0> filter = Builders<MongoBlobIndexModelV0>.Filter.Where(m => m.Ns == ns.ToString() && m.BlobId == id.ToString());
|
|
await collection.FindOneAndReplaceAsync(filter, model, new FindOneAndReplaceOptions<MongoBlobIndexModelV0, MongoBlobIndexModelV0>
|
|
{
|
|
IsUpsert = true
|
|
}, cancellationToken);
|
|
|
|
}
|
|
|
|
public async Task RemoveBlobFromAllRegionsAsync(NamespaceId ns, BlobId id, CancellationToken cancellationToken = default)
|
|
{
|
|
IMongoCollection<MongoBlobIndexModelV0> collection = GetCollection<MongoBlobIndexModelV0>();
|
|
FilterDefinition<MongoBlobIndexModelV0> filter = Builders<MongoBlobIndexModelV0>.Filter.Where(m => m.Ns == ns.ToString() && m.BlobId == id.ToString());
|
|
await collection.DeleteOneAsync(filter, cancellationToken);
|
|
}
|
|
|
|
public IAsyncEnumerable<(NamespaceId, BaseBlobReference)> GetAllBlobReferencesAsync(CancellationToken cancellationToken = default)
|
|
{
|
|
throw new NotImplementedException();
|
|
}
|
|
|
|
public async IAsyncEnumerable<BaseBlobReference> GetBlobReferencesAsync(NamespaceId ns, BlobId id, [EnumeratorCancellation] CancellationToken cancellationToken)
|
|
{
|
|
MongoBlobIndexModelV0? blobInfo = await GetBlobInfoAsync(ns, id, cancellationToken);
|
|
if (blobInfo == null)
|
|
{
|
|
yield break;
|
|
}
|
|
|
|
foreach (Dictionary<string, string> reference in blobInfo.References)
|
|
{
|
|
if (reference.ContainsKey("bucket"))
|
|
{
|
|
string bucket = reference["bucket"];
|
|
string key = reference["key"];
|
|
yield return new RefBlobReference(id, new BucketId(bucket), new RefId(key));
|
|
}
|
|
else if (reference.ContainsKey("blob_id"))
|
|
{
|
|
string blobId = reference["blob_id"];
|
|
yield return new BlobToBlobReference(id, new BlobId(blobId));
|
|
}
|
|
else
|
|
{
|
|
throw new NotImplementedException();
|
|
}
|
|
}
|
|
}
|
|
|
|
public async Task<bool> BlobExistsInRegionAsync(NamespaceId ns, BlobId blobIdentifier, string? region = null, CancellationToken cancellationToken = default)
|
|
{
|
|
MongoBlobIndexModelV0? blobInfo = await GetBlobInfoAsync(ns, blobIdentifier, cancellationToken);
|
|
return blobInfo?.Regions.Contains(_jupiterSettings.CurrentValue.CurrentSite) ?? false;
|
|
}
|
|
|
|
public async Task AddRefToBlobsAsync(NamespaceId ns, BucketId bucket, RefId key, BlobId[] blobs, CancellationToken cancellationToken)
|
|
{
|
|
IMongoCollection<MongoBlobIndexModelV0> collection = GetCollection<MongoBlobIndexModelV0>();
|
|
|
|
string nsAsString = ns.ToString();
|
|
Task[] refUpdateTasks = new Task[blobs.Length];
|
|
for (int i = 0; i < blobs.Length; i++)
|
|
{
|
|
BlobId id = blobs[i];
|
|
refUpdateTasks[i] = Task.Run(async () =>
|
|
{
|
|
UpdateDefinition<MongoBlobIndexModelV0> update = Builders<MongoBlobIndexModelV0>.Update.AddToSet(m => m.References, new Dictionary<string, string> { { "bucket", bucket.ToString() }, { "key", key.ToString() } });
|
|
FilterDefinition<MongoBlobIndexModelV0> filter = Builders<MongoBlobIndexModelV0>.Filter.Where(m => m.Ns == nsAsString && m.BlobId == id.ToString());
|
|
|
|
await collection.FindOneAndUpdateAsync(filter, update, cancellationToken: cancellationToken);
|
|
}, cancellationToken);
|
|
}
|
|
|
|
await Task.WhenAll(refUpdateTasks);
|
|
}
|
|
|
|
public async IAsyncEnumerable<(NamespaceId, BlobId)> GetAllBlobsAsync([EnumeratorCancellation] CancellationToken cancellationToken)
|
|
{
|
|
IMongoCollection<MongoBlobIndexModelV0> collection = GetCollection<MongoBlobIndexModelV0>();
|
|
IAsyncCursor<MongoBlobIndexModelV0>? cursor = await collection.FindAsync(FilterDefinition<MongoBlobIndexModelV0>.Empty, cancellationToken: cancellationToken);
|
|
|
|
while (await cursor.MoveNextAsync(cancellationToken))
|
|
{
|
|
foreach (MongoBlobIndexModelV0 model in cursor.Current)
|
|
{
|
|
yield return (new NamespaceId(model.Ns), new BlobId(model.BlobId));
|
|
}
|
|
}
|
|
}
|
|
|
|
public async Task RemoveReferencesAsync(NamespaceId ns, BlobId id, List<BaseBlobReference>? referencesToRemove, CancellationToken cancellationToken)
|
|
{
|
|
IMongoCollection<MongoBlobIndexModelV0> collection = GetCollection<MongoBlobIndexModelV0>();
|
|
|
|
string nsAsString = ns.ToString();
|
|
if (referencesToRemove == null)
|
|
{
|
|
FilterDefinition<MongoBlobIndexModelV0> filter = Builders<MongoBlobIndexModelV0>.Filter.Where(m => m.Ns == nsAsString && m.BlobId == id.ToString());
|
|
await collection.DeleteOneAsync(filter, cancellationToken);
|
|
}
|
|
else
|
|
{
|
|
List<Dictionary<string, string>> refs = referencesToRemove.Select(reference =>
|
|
{
|
|
if (reference is RefBlobReference refBlobReference)
|
|
{
|
|
return new Dictionary<string, string>
|
|
{
|
|
{ "bucket", refBlobReference.Bucket.ToString() }, { "key", refBlobReference.Key.ToString() }
|
|
};
|
|
}
|
|
else if (reference is BlobToBlobReference blobToBlobReference)
|
|
{
|
|
return new Dictionary<string, string>
|
|
{
|
|
{ "blob_id", blobToBlobReference.Blob.ToString()}
|
|
};
|
|
}
|
|
else
|
|
{
|
|
throw new NotImplementedException();
|
|
}
|
|
}).ToList();
|
|
UpdateDefinition<MongoBlobIndexModelV0> update = Builders<MongoBlobIndexModelV0>.Update.PullAll(m => m.References, refs);
|
|
FilterDefinition<MongoBlobIndexModelV0> filter = Builders<MongoBlobIndexModelV0>.Filter.Where(m => m.Ns == nsAsString && m.BlobId == id.ToString());
|
|
|
|
await collection.FindOneAndUpdateAsync(filter, update, cancellationToken: cancellationToken);
|
|
}
|
|
}
|
|
|
|
public async Task<List<string>> GetBlobRegionsAsync(NamespaceId ns, BlobId blob, CancellationToken cancellationToken)
|
|
{
|
|
MongoBlobIndexModelV0? blobInfo = await GetBlobInfoAsync(ns, blob, cancellationToken);
|
|
if (blobInfo == null)
|
|
{
|
|
throw new BlobNotFoundException(ns, blob);
|
|
}
|
|
return blobInfo.Regions.ToList();
|
|
}
|
|
|
|
public async Task AddBlobReferencesAsync(NamespaceId ns, BlobId sourceBlob, BlobId targetBlob, CancellationToken cancellationToken)
|
|
{
|
|
IMongoCollection<MongoBlobIndexModelV0> collection = GetCollection<MongoBlobIndexModelV0>();
|
|
|
|
string nsAsString = ns.ToString();
|
|
|
|
UpdateDefinition<MongoBlobIndexModelV0> update = Builders<MongoBlobIndexModelV0>.Update.AddToSet(m => m.References, new Dictionary<string, string> { { "blob_id", targetBlob.ToString() } });
|
|
FilterDefinition<MongoBlobIndexModelV0> filter = Builders<MongoBlobIndexModelV0>.Filter.Where(m => m.Ns == nsAsString && m.BlobId == sourceBlob.ToString());
|
|
|
|
await collection.FindOneAndUpdateAsync(filter, update, cancellationToken: cancellationToken);
|
|
}
|
|
|
|
public async Task AddBlobToBucketListAsync(NamespaceId ns, BucketId bucket, RefId key, BlobId blobId, long blobSize, CancellationToken cancellationToken)
|
|
{
|
|
IMongoCollection<MongoBucketBlobV0> collection = GetCollection<MongoBucketBlobV0>();
|
|
|
|
FilterDefinition<MongoBucketBlobV0> filter = Builders<MongoBucketBlobV0>.Filter.Where(m => m.Ns == ns.ToString() && m.BucketId == bucket.ToString() && m.RefId == key.ToString() && m.BlobId == blobId.ToString());
|
|
await collection.ReplaceOneAsync(filter, new MongoBucketBlobV0(ns, bucket, key, blobId, blobSize), new ReplaceOptions() { IsUpsert = true }, cancellationToken);
|
|
}
|
|
|
|
public async Task RemoveBlobFromBucketListAsync(NamespaceId ns, BucketId bucket, RefId key, List<BlobId> blobIds, CancellationToken cancellationToken)
|
|
{
|
|
IMongoCollection<MongoBucketBlobV0> collection = GetCollection<MongoBucketBlobV0>();
|
|
|
|
FilterDefinition<MongoBucketBlobV0> filter = Builders<MongoBucketBlobV0>.Filter.Where(m => m.Ns == ns.ToString() && m.BucketId == bucket.ToString() && m.RefId == key.ToString());
|
|
await collection.DeleteManyAsync(filter, cancellationToken);
|
|
}
|
|
|
|
public async Task<BucketStats> CalculateBucketStatisticsAsync(NamespaceId ns, BucketId bucket, CancellationToken cancellationToken)
|
|
{
|
|
IMongoCollection<MongoBucketBlobV0> collection = GetCollection<MongoBucketBlobV0>();
|
|
FilterDefinition<MongoBucketBlobV0> filter = Builders<MongoBucketBlobV0>.Filter.Where(m => m.Ns == ns.ToString() && m.BucketId == bucket.ToString());
|
|
|
|
var blobStats = await collection.Aggregate()
|
|
.Match(filter)
|
|
.Group(
|
|
a => a.BucketId,
|
|
r => new
|
|
{
|
|
TotalSize = r.Sum(a => a.Size),
|
|
SmallestBlob = r.Min(a => a.Size),
|
|
LargestBlob = r.Max(a => a.Size),
|
|
CountOfBlobs = r.Count(),
|
|
}).ToListAsync(cancellationToken);
|
|
|
|
var blobStat = blobStats.FirstOrDefault();
|
|
|
|
long countOfRefs = collection
|
|
.Aggregate()
|
|
.Match(filter)
|
|
.Group(m => m.RefId,
|
|
grouping => new { DoesNotMatter = grouping.Key })
|
|
.Count()
|
|
.First(cancellationToken)
|
|
.Count;
|
|
|
|
return new BucketStats()
|
|
{
|
|
Namespace = ns,
|
|
Bucket = bucket,
|
|
CountOfRefs = countOfRefs,
|
|
CountOfBlobs = blobStat?.CountOfBlobs ?? 0,
|
|
SmallestBlobFound = (long)(blobStat?.SmallestBlob ?? 0),
|
|
LargestBlob = (long)(blobStat?.LargestBlob ?? 0),
|
|
TotalSize = blobStat?.TotalSize ?? 0,
|
|
AvgSize = blobStat == null ? 0 : blobStat.TotalSize / (double)blobStat.CountOfBlobs
|
|
};
|
|
}
|
|
}
|
|
|
|
[BsonDiscriminator("blob-index.v0")]
|
|
[BsonIgnoreExtraElements]
|
|
[MongoCollectionName("BlobIndex")]
|
|
class MongoBlobIndexModelV0
|
|
{
|
|
[BsonConstructor]
|
|
public MongoBlobIndexModelV0(string ns, string blobId, List<string> regions, List<Dictionary<string, string>> references)
|
|
{
|
|
Ns = ns;
|
|
BlobId = blobId;
|
|
Regions = regions;
|
|
References = references;
|
|
}
|
|
|
|
public MongoBlobIndexModelV0(NamespaceId ns, BlobId blobId)
|
|
{
|
|
Ns = ns.ToString();
|
|
BlobId = blobId.ToString();
|
|
}
|
|
|
|
[BsonRequired]
|
|
public string Ns { get; set; }
|
|
|
|
[BsonRequired]
|
|
public string BlobId { get; set; }
|
|
|
|
public List<string> Regions { get; set; } = new List<string>();
|
|
|
|
[BsonDictionaryOptions(DictionaryRepresentation.Document)]
|
|
public List<Dictionary<string, string>> References { get; set; } = new List<Dictionary<string, string>>();
|
|
|
|
}
|
|
|
|
[BsonDiscriminator("bucket-blob.v0")]
|
|
[BsonIgnoreExtraElements]
|
|
[MongoCollectionName("BucketBlob")]
|
|
class MongoBucketBlobV0
|
|
{
|
|
[BsonConstructor]
|
|
public MongoBucketBlobV0(string ns, string bucketId, string refId, string blobId, long size)
|
|
{
|
|
Ns = ns;
|
|
BucketId = bucketId;
|
|
RefId = refId;
|
|
BlobId = blobId;
|
|
Size = size;
|
|
}
|
|
|
|
public MongoBucketBlobV0(NamespaceId ns, BucketId bucketId, RefId refId, BlobId blobId, long size)
|
|
{
|
|
Ns = ns.ToString();
|
|
BucketId = bucketId.ToString();
|
|
RefId = refId.ToString();
|
|
BlobId = blobId.ToString();
|
|
Size = size;
|
|
}
|
|
|
|
[BsonRequired]
|
|
public string Ns { get; set; }
|
|
|
|
[BsonRequired]
|
|
public string BucketId { get; set; }
|
|
|
|
[BsonRequired]
|
|
public string RefId { get; set; }
|
|
|
|
[BsonRequired]
|
|
public string BlobId { get; set; }
|
|
|
|
[BsonRequired]
|
|
public long Size { get; set; }
|
|
} |