|
// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
using System.Collections.Concurrent;
using System.Globalization;
using System.Text.Json.Nodes;
namespace Aspire.Hosting.RemoteHost.Ats;
/// <summary>
/// Manages registration and lookup of ATS handles for capability dispatch.
/// Handles are opaque typed references with IDs in the format: {typeId}:{instanceId}
/// </summary>
internal sealed class HandleRegistry : IAsyncDisposable
{
private readonly ConcurrentDictionary<string, HandleEntry> _handles = new();
private long _idCounter;
/// <summary>
/// Represents a registered handle entry.
/// </summary>
private sealed class HandleEntry
{
public required object Object { get; init; }
public required string TypeId { get; init; }
public required long InstanceId { get; init; }
}
/// <summary>
/// Registers an object as a handle with the specified ATS type ID.
/// </summary>
/// <param name="obj">The object to register.</param>
/// <param name="typeId">The ATS type ID (e.g., "aspire.redis/RedisBuilder").</param>
/// <returns>The handle ID (just the instance number).</returns>
public string Register(object obj, string typeId)
{
var instanceId = Interlocked.Increment(ref _idCounter);
var handleId = instanceId.ToString(CultureInfo.InvariantCulture);
_handles[handleId] = new HandleEntry
{
Object = obj,
TypeId = typeId,
InstanceId = instanceId
};
return handleId;
}
/// <summary>
/// Tries to get a handle entry by its ID.
/// </summary>
/// <param name="handleId">The handle ID.</param>
/// <param name="obj">The underlying object if found.</param>
/// <param name="typeId">The ATS type ID if found.</param>
/// <returns>True if the handle was found, false otherwise.</returns>
public bool TryGet(string handleId, out object? obj, out string? typeId)
{
if (_handles.TryGetValue(handleId, out var entry))
{
obj = entry.Object;
typeId = entry.TypeId;
return true;
}
obj = null;
typeId = null;
return false;
}
/// <summary>
/// Gets the underlying object for a handle.
/// </summary>
/// <param name="handleId">The handle ID.</param>
/// <returns>The underlying object.</returns>
/// <exception cref="InvalidOperationException">Thrown if the handle is not found.</exception>
public object GetObject(string handleId)
{
if (!_handles.TryGetValue(handleId, out var entry))
{
throw new InvalidOperationException($"Handle '{handleId}' not found in registry");
}
return entry.Object;
}
/// <summary>
/// Gets the underlying object for a handle, cast to the specified type.
/// </summary>
/// <typeparam name="T">The expected type.</typeparam>
/// <param name="handleId">The handle ID.</param>
/// <returns>The underlying object.</returns>
/// <exception cref="InvalidOperationException">Thrown if the handle is not found or type doesn't match.</exception>
public T GetObject<T>(string handleId) where T : class
{
var obj = GetObject(handleId);
if (obj is not T typed)
{
throw new InvalidOperationException(
$"Handle '{handleId}' contains {obj.GetType().FullName}, expected {typeof(T).FullName}");
}
return typed;
}
/// <summary>
/// Gets the ATS type ID for a handle.
/// </summary>
/// <param name="handleId">The handle ID.</param>
/// <returns>The ATS type ID.</returns>
/// <exception cref="InvalidOperationException">Thrown if the handle is not found.</exception>
public string GetTypeId(string handleId)
{
if (!_handles.TryGetValue(handleId, out var entry))
{
throw new InvalidOperationException($"Handle '{handleId}' not found in registry");
}
return entry.TypeId;
}
/// <summary>
/// Checks if a handle exists in the registry.
/// </summary>
/// <param name="handleId">The handle ID.</param>
/// <returns>True if the handle exists, false otherwise.</returns>
public bool Contains(string handleId)
{
return _handles.ContainsKey(handleId);
}
/// <summary>
/// Unregisters a handle from the registry.
/// </summary>
/// <param name="handleId">The handle ID to unregister.</param>
/// <returns>True if the handle was removed, false if it wasn't found.</returns>
public bool Unregister(string handleId)
{
return _handles.TryRemove(handleId, out _);
}
/// <summary>
/// Marshals an object to a JSON handle reference.
/// </summary>
/// <param name="obj">The object to marshal.</param>
/// <param name="typeId">The ATS type ID.</param>
/// <returns>A JsonObject containing the handle reference.</returns>
public JsonObject Marshal(object obj, string typeId)
{
var handleId = Register(obj, typeId);
return new JsonObject
{
["$handle"] = handleId,
["$type"] = typeId
};
}
/// <summary>
/// Gets the count of registered handles.
/// </summary>
public int Count => _handles.Count;
/// <summary>
/// Disposes all disposable objects in the registry and clears it.
/// </summary>
public async ValueTask DisposeAsync()
{
var entries = _handles.Values.ToList();
_handles.Clear();
foreach (var entry in entries)
{
if (entry.Object is IAsyncDisposable asyncDisposable)
{
await asyncDisposable.DisposeAsync().ConfigureAwait(false);
}
else if (entry.Object is IDisposable disposable)
{
disposable.Dispose();
}
}
}
}
/// <summary>
/// Reference to an ATS handle. Used when passing handles as arguments.
/// JSON shape: { "$handle": "42", "$type": "Aspire.Hosting.Redis/..." }
/// </summary>
internal sealed class HandleRef
{
/// <summary>
/// The handle identifier.
/// </summary>
public required string HandleId { get; init; }
/// <summary>
/// Creates a HandleRef from a JSON node if it contains a $handle property.
/// </summary>
public static HandleRef? FromJsonNode(JsonNode? node)
{
if (node is JsonObject obj && obj.TryGetPropertyValue("$handle", out var handleNode))
{
var handleId = handleNode?.GetValue<string>();
if (!string.IsNullOrEmpty(handleId))
{
return new HandleRef { HandleId = handleId };
}
}
return null;
}
/// <summary>
/// Checks if a JSON node is a handle reference.
/// </summary>
public static bool IsHandleRef(JsonNode? node)
{
return node is JsonObject obj && obj.ContainsKey("$handle");
}
}
|