Skip to content

Commit 870a911

Browse files
adityamandaleekaadityamandaleeka
committed
Make CosmosDB initialization be lazy.
1 parent b80426e commit 870a911

File tree

4 files changed

+108
-70
lines changed

4 files changed

+108
-70
lines changed

‎dotnet/samples/HelloHttpApi/HelloHttpApi.ApiService/HostApplicationBuilderAgentExtensions.cs‎

Lines changed: 2 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,6 @@
99

1010
namespace HelloHttpApi.ApiService;
1111

12-
#pragma warning disable VSTHRD002
13-
1412
public static class HostApplicationBuilderAgentExtensions
1513
{
1614
public static IHostApplicationBuilder AddAIAgent(this IHostApplicationBuilder builder, string name, string instructions, string? chatClientKey = null)
@@ -34,22 +32,10 @@ public static IHostApplicationBuilder AddAIAgent(this IHostApplicationBuilder bu
3432
var actorBuilder = builder.AddActorRuntime();
3533

3634
// Add CosmosDB state storage to override default storage
37-
builder.Services.AddSingleton<IActorStateStorage>(serviceProvider =>
35+
builder.Services.AddCosmosActorStateStorage(serviceProvider =>
3836
{
3937
var cosmosClient = serviceProvider.GetRequiredService<CosmosClient>();
40-
41-
// Use the database from Aspire configuration
42-
var database = cosmosClient.GetDatabase("actor-state-db");
43-
44-
// Create container if it doesn't exist with proper configuration
45-
var containerProperties = new ContainerProperties
46-
{
47-
Id = "ActorState",
48-
PartitionKeyPath = "/actorId"
49-
};
50-
51-
var container = database.CreateContainerIfNotExistsAsync(containerProperties).GetAwaiter().GetResult();
52-
return new CosmosActorStateStorage(container.Container);
38+
return new LazyCosmosContainer(cosmosClient, "actor-state-db", "ActorState");
5339
});
5440

5541
actorBuilder.AddActorType(

‎dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/CosmosActorStateStorage.cs‎

Lines changed: 19 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -35,15 +35,22 @@ namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB;
3535
/// </summary>
3636
public class CosmosActorStateStorage : IActorStateStorage
3737
{
38-
private readonly Container _container;
38+
private readonly LazyCosmosContainer _lazyContainer;
3939
private const string InitialEtag = "0"; // Initial ETag value when no state exists
4040

4141
/// <summary>
4242
/// Constructs a new instance of <see cref="CosmosActorStateStorage"/> with the specified Cosmos DB container.
4343
/// </summary>
4444
/// <param name="container">The Cosmos DB container to use for storage.</param>
4545
/// <exception cref="ArgumentNullException">Thrown when <paramref name="container"/> is null.</exception>
46-
public CosmosActorStateStorage(Container container) => this._container = container ?? throw new ArgumentNullException(nameof(container));
46+
public CosmosActorStateStorage(Container container) => this._lazyContainer = new LazyCosmosContainer(container);
47+
48+
/// <summary>
49+
/// Constructs a new instance of <see cref="CosmosActorStateStorage"/> with the specified lazy container.
50+
/// </summary>
51+
/// <param name="lazyContainer">The lazy Cosmos DB container to use for storage.</param>
52+
/// <exception cref="ArgumentNullException">Thrown when <paramref name="lazyContainer"/> is null.</exception>
53+
public CosmosActorStateStorage(LazyCosmosContainer lazyContainer) => this._lazyContainer = lazyContainer ?? throw new ArgumentNullException(nameof(lazyContainer));
4754

4855
/// <summary>
4956
/// Writes state changes to the actor's persistent storage.
@@ -54,6 +61,8 @@ public async ValueTask<WriteResponse> WriteStateAsync(
5461
string etag,
5562
CancellationToken cancellationToken = default)
5663
{
64+
var container = await this._lazyContainer.GetContainerAsync().ConfigureAwait(false);
65+
5766
if (operations.Count == 0)
5867
{
5968
// No operations to perform - return success with current ETag or generate new one
@@ -62,14 +71,14 @@ public async ValueTask<WriteResponse> WriteStateAsync(
6271
}
6372

6473
var partitionKey = new PartitionKey(actorId.ToString());
65-
var batch = this._container.CreateTransactionalBatch(partitionKey);
74+
var batch = container.CreateTransactionalBatch(partitionKey);
6675

6776
// First, try to read existing root document to get current version
6877
var rootDocId = GetRootDocumentId(actorId);
6978
ActorRootDocument? existingRoot = null;
7079
try
7180
{
72-
var rootResponse = await this._container.ReadItemAsync<ActorRootDocument>(
81+
var rootResponse = await container.ReadItemAsync<ActorRootDocument>(
7382
rootDocId,
7483
partitionKey,
7584
cancellationToken: cancellationToken).ConfigureAwait(false);
@@ -165,10 +174,11 @@ public async ValueTask<ReadResponse> ReadStateAsync(
165174
IReadOnlyCollection<ActorStateReadOperation> operations,
166175
CancellationToken cancellationToken = default)
167176
{
177+
var container = await this._lazyContainer.GetContainerAsync().ConfigureAwait(false);
168178
var results = new List<ActorReadResult>();
169179

170180
// Read root document first to get actor-level ETag
171-
string actorETag = await this.GetActorETagAsync(actorId, cancellationToken).ConfigureAwait(false);
181+
string actorETag = await this.GetActorETagAsync(container, actorId, cancellationToken).ConfigureAwait(false);
172182

173183
foreach (var op in operations)
174184
{
@@ -178,7 +188,7 @@ public async ValueTask<ReadResponse> ReadStateAsync(
178188
var id = GetDocumentId(actorId, get.Key);
179189
try
180190
{
181-
var response = await this._container.ReadItemAsync<ActorStateDocument>(
191+
var response = await container.ReadItemAsync<ActorStateDocument>(
182192
id,
183193
new PartitionKey(actorId.ToString()),
184194
cancellationToken: cancellationToken)
@@ -215,7 +225,7 @@ public async ValueTask<ReadResponse> ReadStateAsync(
215225
MaxItemCount = 100 // TODO Fix
216226
};
217227

218-
var iterator = this._container.GetItemQueryIterator<KeyProjection>(
228+
var iterator = container.GetItemQueryIterator<KeyProjection>(
219229
query,
220230
list.ContinuationToken,
221231
requestOptions);
@@ -299,12 +309,12 @@ private static string GetRootDocumentId(ActorId actorId)
299309
/// Gets the current ETag for the actor's root document.
300310
/// Returns a generated ETag if no root document exists.
301311
/// </summary>
302-
private async ValueTask<string> GetActorETagAsync(ActorId actorId, CancellationToken cancellationToken)
312+
private async ValueTask<string> GetActorETagAsync(Container container, ActorId actorId, CancellationToken cancellationToken)
303313
{
304314
var rootDocId = GetRootDocumentId(actorId);
305315
try
306316
{
307-
var rootResponse = await this._container.ReadItemAsync<ActorRootDocument>(
317+
var rootResponse = await container.ReadItemAsync<ActorRootDocument>(
308318
rootDocId,
309319
new PartitionKey(actorId.ToString()),
310320
cancellationToken: cancellationToken).ConfigureAwait(false);

‎dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/LazyCosmosContainer.cs‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
// Copyright (c) Microsoft. All rights reserved.
2+
3+
using System;
4+
using System.Collections.ObjectModel;
5+
using System.Threading;
6+
using System.Threading.Tasks;
7+
using Microsoft.Azure.Cosmos;
8+
9+
namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB;
10+
11+
#pragma warning disable VSTHRD011 // Use AsyncLazy<T>
12+
13+
/// <summary>
14+
/// A lazy wrapper around a Cosmos DB Container.
15+
/// This avoids performing async I/O-bound operations (i.e. Cosmos DB setup) during
16+
/// DI registration, deferring them until first access.
17+
/// </summary>
18+
public sealed class LazyCosmosContainer
19+
{
20+
private readonly CosmosClient? _cosmosClient;
21+
private readonly string? _databaseName;
22+
private readonly string? _containerName;
23+
private readonly Lazy<Task<Container>> _lazyContainer;
24+
25+
/// <summary>
26+
/// LazyCosmosContainer constructor that initializes the container lazily.
27+
/// </summary>
28+
public LazyCosmosContainer(CosmosClient cosmosClient, string databaseName, string containerName)
29+
{
30+
this._cosmosClient = cosmosClient ?? throw new ArgumentNullException(nameof(cosmosClient));
31+
this._databaseName = databaseName ?? throw new ArgumentNullException(nameof(databaseName));
32+
this._containerName = containerName ?? throw new ArgumentNullException(nameof(containerName));
33+
this._lazyContainer = new Lazy<Task<Container>>(this.InitializeContainerAsync, LazyThreadSafetyMode.ExecutionAndPublication);
34+
}
35+
36+
/// <summary>
37+
/// LazyCosmosContainer constructor that accepts an existing Container instance.
38+
/// </summary>
39+
public LazyCosmosContainer(Container container)
40+
{
41+
if (container is null)
42+
{
43+
throw new ArgumentNullException(nameof(container));
44+
}
45+
46+
this._lazyContainer = new Lazy<Task<Container>>(() => Task.FromResult(container), LazyThreadSafetyMode.ExecutionAndPublication);
47+
}
48+
49+
/// <summary>
50+
/// Gets the Container, initializing it if necessary.
51+
/// </summary>
52+
public Task<Container> GetContainerAsync() => this._lazyContainer.Value;
53+
54+
private async Task<Container> InitializeContainerAsync()
55+
{
56+
// Create database if it doesn't exist
57+
var database = await this._cosmosClient!.CreateDatabaseIfNotExistsAsync(this._databaseName!).ConfigureAwait(false);
58+
59+
var containerProperties = new ContainerProperties(this._containerName!, "/actorId")
60+
{
61+
Id = this._containerName!,
62+
IndexingPolicy = new IndexingPolicy
63+
{
64+
IndexingMode = IndexingMode.Consistent,
65+
Automatic = true
66+
},
67+
PartitionKeyPaths = ["/actorId"]
68+
};
69+
70+
// Add composite index for efficient queries
71+
containerProperties.IndexingPolicy.CompositeIndexes.Add(new Collection<CompositePath>
72+
{
73+
new() { Path = "/actorId", Order = CompositePathSortOrder.Ascending },
74+
new() { Path = "/key", Order = CompositePathSortOrder.Ascending }
75+
});
76+
77+
var container = await database.Database.CreateContainerIfNotExistsAsync(containerProperties).ConfigureAwait(false);
78+
return container.Container;
79+
}
80+
}

‎dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/ServiceCollectionExtensions.cs‎

Lines changed: 7 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
// Copyright (c) Microsoft. All rights reserved.
22

33
using System;
4-
using System.Collections.ObjectModel;
54
using System.Text.Json;
65
using Microsoft.Azure.Cosmos;
76
using Microsoft.Extensions.DependencyInjection;
@@ -46,34 +45,11 @@ public static IServiceCollection AddCosmosActorStateStorage(
4645
return new CosmosClient(connectionString, cosmosClientOptions);
4746
});
4847

49-
// Register Container as singleton
50-
services.AddSingleton<Container>(serviceProvider =>
48+
// Register LazyCosmosContainer as singleton
49+
services.AddSingleton<LazyCosmosContainer>(serviceProvider =>
5150
{
5251
var cosmosClient = serviceProvider.GetRequiredService<CosmosClient>();
53-
54-
// Create database and container if they don't exist
55-
var database = cosmosClient.CreateDatabaseIfNotExistsAsync(databaseName).GetAwaiter().GetResult();
56-
57-
var containerProperties = new ContainerProperties(containerName, "/actorId")
58-
{
59-
Id = containerName,
60-
IndexingPolicy = new IndexingPolicy
61-
{
62-
IndexingMode = IndexingMode.Consistent,
63-
Automatic = true
64-
},
65-
PartitionKeyPaths = ["/actorId"]
66-
};
67-
68-
// Add composite index for efficient queries
69-
containerProperties.IndexingPolicy.CompositeIndexes.Add(new Collection<CompositePath>
70-
{
71-
new() { Path = "/actorId", Order = CompositePathSortOrder.Ascending },
72-
new() { Path = "/key", Order = CompositePathSortOrder.Ascending }
73-
});
74-
75-
var container = database.Database.CreateContainerIfNotExistsAsync(containerProperties).GetAwaiter().GetResult();
76-
return container.Container;
52+
return new LazyCosmosContainer(cosmosClient, databaseName, containerName);
7753
});
7854

7955
// Register the storage implementation
@@ -83,30 +59,16 @@ public static IServiceCollection AddCosmosActorStateStorage(
8359
}
8460

8561
/// <summary>
86-
/// Adds Cosmos DB actor state storage to the service collection with a factory for the Container.
62+
/// Adds Cosmos DB actor state storage to the service collection with a factory for the LazyCosmosContainer.
8763
/// </summary>
8864
/// <param name="services">The service collection to add services to.</param>
89-
/// <param name="containerFactory">A factory function that creates the Cosmos Container.</param>
65+
/// <param name="lazyContainerFactory">A factory function that creates the LazyCosmosContainer.</param>
9066
/// <returns>The service collection for chaining.</returns>
9167
public static IServiceCollection AddCosmosActorStateStorage(
9268
this IServiceCollection services,
93-
Func<IServiceProvider, Container> containerFactory)
94-
{
95-
services.AddSingleton<Container>(containerFactory);
96-
services.AddSingleton<IActorStateStorage, CosmosActorStateStorage>();
97-
return services;
98-
}
99-
100-
/// <summary>
101-
/// Adds Cosmos DB actor state storage to the service collection using an existing Container registration.
102-
/// </summary>
103-
/// <param name="services">The service collection to add services to.</param>
104-
/// <returns>The service collection for chaining.</returns>
105-
/// <remarks>
106-
/// This overload assumes that a <see cref="Container"/> is already registered in the service collection.
107-
/// </remarks>
108-
public static IServiceCollection AddCosmosActorStateStorage(this IServiceCollection services)
69+
Func<IServiceProvider, LazyCosmosContainer> lazyContainerFactory)
10970
{
71+
services.AddSingleton<LazyCosmosContainer>(lazyContainerFactory);
11072
services.AddSingleton<IActorStateStorage, CosmosActorStateStorage>();
11173
return services;
11274
}

0 commit comments

Comments
 (0)