From e5d0ae67602686489d95eb957e44148264934e66 Mon Sep 17 00:00:00 2001 From: Sophia Tevosyan Date: Mon, 30 Jun 2025 14:55:13 -0700 Subject: [PATCH] first commit --- src/Shared/Grpc/ProtoUtils.cs | 14 +- src/Worker/Core/DurableTaskWorkerOptions.cs | 19 +- .../Grpc/GrpcDurableTaskWorker.Processor.cs | 243 ++++++++++++++---- src/Worker/Grpc/GrpcDurableTaskWorker.cs | 3 +- src/Worker/Grpc/GrpcOrchestrationRunner.cs | 5 +- src/Worker/Grpc/LRUCache.cs | 222 ++++++++++++++++ 6 files changed, 447 insertions(+), 59 deletions(-) create mode 100644 src/Worker/Grpc/LRUCache.cs diff --git a/src/Shared/Grpc/ProtoUtils.cs b/src/Shared/Grpc/ProtoUtils.cs index 868ecc66..4e8eefc6 100644 --- a/src/Shared/Grpc/ProtoUtils.cs +++ b/src/Shared/Grpc/ProtoUtils.cs @@ -275,6 +275,8 @@ internal static Timestamp ToTimestamp(this DateTime dateTime) /// value that was provided by the corresponding that triggered the orchestrator execution. /// /// The entity conversion state, or null if no conversion is required. + /// Upon completion of the method, contains the status of this orchestration. + /// Whether or not the worker has cached the orchestration history. False by default. /// The orchestrator response. /// When an orchestrator action is unknown. internal static P.OrchestratorResponse ConstructOrchestratorResponse( @@ -282,7 +284,9 @@ internal static P.OrchestratorResponse ConstructOrchestratorResponse( string? customStatus, IEnumerable actions, string completionToken, - EntityConversionState? entityConversionState) + EntityConversionState? entityConversionState, + out OrchestrationStatus orchestrationStatus, + bool historyCached = false) { Check.NotNull(actions); var response = new P.OrchestratorResponse @@ -290,7 +294,9 @@ internal static P.OrchestratorResponse ConstructOrchestratorResponse( InstanceId = instanceId, CustomStatus = customStatus, CompletionToken = completionToken, + HistoryCached = historyCached, }; + orchestrationStatus = OrchestrationStatus.Running; // default status foreach (OrchestratorAction action in actions) { @@ -417,6 +423,7 @@ internal static P.OrchestratorResponse ConstructOrchestratorResponse( OrchestrationStatus = completeAction.OrchestrationStatus.ToProtobuf(), Result = completeAction.Result, }; + orchestrationStatus = completeAction.OrchestrationStatus; if (completeAction.OrchestrationStatus == OrchestrationStatus.Failed) { @@ -822,12 +829,14 @@ internal static void ToEntityBatchRequest( /// The operation result to convert. /// The completion token, or null for the older protocol. /// Additional information about each operation, required by DTS. + /// Indicates whether the worker has cached the entity state. /// The converted operation result. [return: NotNullIfNotNull(nameof(entityBatchResult))] internal static P.EntityBatchResult? ToEntityBatchResult( this EntityBatchResult? entityBatchResult, string? completionToken = null, - IEnumerable? operationInfos = null) + IEnumerable? operationInfos = null, + bool entityStateCached = false) { if (entityBatchResult == null) { @@ -842,6 +851,7 @@ internal static void ToEntityBatchRequest( Results = { entityBatchResult.Results?.Select(a => a.ToOperationResult()) ?? [] }, CompletionToken = completionToken ?? string.Empty, OperationInfos = { operationInfos ?? [] }, + EntityStateCached = entityStateCached, }; } diff --git a/src/Worker/Core/DurableTaskWorkerOptions.cs b/src/Worker/Core/DurableTaskWorkerOptions.cs index b2bdca04..e71969f0 100644 --- a/src/Worker/Core/DurableTaskWorkerOptions.cs +++ b/src/Worker/Core/DurableTaskWorkerOptions.cs @@ -143,7 +143,22 @@ public DataConverter DataConverter /// /// Gets a value indicating whether versioning is explicitly set or not. /// - public bool IsVersioningSet { get; internal set; } + public bool IsVersioningSet { get; internal set; } + + /// + /// Gets or sets the maximum size of the orchestration history cache in bytes. + /// + public int OrchestrationHistoryCacheSizeInBytes { get; set; } = 10000; + + /// + /// Gets or sets the amount of time that an orchestration history can remain in the cache before it is considered "stale" and evicted. + /// + public int OrchestrationHistoryCacheStaleEvictionTimeInMilliseconds { get; set; } = 30 * 1000; // 30 seconds + + /// + /// Gets or sets how often the cache will check for stale items and evict them. + /// + public int OrchestrationHistoryCacheCheckForStaleItemsPeriodInMilliseconds { get; set; } = 5 * 1000; // 5 seconds /// /// Gets a value indicating whether was explicitly set or not. @@ -154,7 +169,7 @@ public DataConverter DataConverter /// will not resolve it. If not set, we will attempt to resolve it. This is so the /// behavior is consistently irrespective of option configuration ordering. /// - internal bool DataConverterExplicitlySet { get; private set; } + internal bool DataConverterExplicitlySet { get; private set; } /// /// Applies these option values to another. diff --git a/src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs b/src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs index 327bb0cc..e4b712b9 100644 --- a/src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs +++ b/src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs @@ -1,11 +1,13 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. - -using System.Text; + +using System; +using System.Text; using DurableTask.Core; using DurableTask.Core.Entities; using DurableTask.Core.Entities.OperationFormat; using DurableTask.Core.History; +using Google.Protobuf.Collections; using Microsoft.DurableTask.Abstractions; using Microsoft.DurableTask.Entities; using Microsoft.DurableTask.Worker.Shims; @@ -22,21 +24,31 @@ namespace Microsoft.DurableTask.Worker.Grpc; /// sealed partial class GrpcDurableTaskWorker { - class Processor + class Processor : IDisposable { static readonly Google.Protobuf.WellKnownTypes.Empty EmptyMessage = new(); readonly GrpcDurableTaskWorker worker; readonly TaskHubSidecarServiceClient client; readonly DurableTaskShimFactory shimFactory; - readonly GrpcDurableTaskWorkerOptions.InternalOptions internalOptions; + readonly GrpcDurableTaskWorkerOptions.InternalOptions internalOptions; + readonly LRUCache> orchestrationHistories; + readonly LRUCache entityStates; public Processor(GrpcDurableTaskWorker worker, TaskHubSidecarServiceClient client) { this.worker = worker; this.client = client; this.shimFactory = new DurableTaskShimFactory(this.worker.grpcOptions, this.worker.loggerFactory); - this.internalOptions = this.worker.grpcOptions.Internal; + this.internalOptions = this.worker.grpcOptions.Internal; + this.orchestrationHistories = new LRUCache>( + worker.workerOptions.OrchestrationHistoryCacheSizeInBytes, + worker.workerOptions.OrchestrationHistoryCacheCheckForStaleItemsPeriodInMilliseconds, + worker.workerOptions.OrchestrationHistoryCacheStaleEvictionTimeInMilliseconds); + this.entityStates = new LRUCache( + worker.workerOptions.OrchestrationHistoryCacheSizeInBytes, + worker.workerOptions.OrchestrationHistoryCacheCheckForStaleItemsPeriodInMilliseconds, + worker.workerOptions.OrchestrationHistoryCacheStaleEvictionTimeInMilliseconds); } ILogger Logger => this.worker.logger; @@ -98,6 +110,14 @@ public async Task ExecuteAsync(CancellationToken cancellation) } } } + + public void Dispose() + { + this.orchestrationHistories.Dispose(); + this.entityStates.Dispose(); + + GC.SuppressFinalize(this); + } static string GetActionsListForLogging(IReadOnlyList actions) { @@ -172,6 +192,25 @@ static string GetActionsListForLogging(IReadOnlyList actio } return failureDetails; + } + + static int CalculateSizeInBytes(RepeatedField historyEvents) + { + int sizeInBytes = 0; + foreach (P.HistoryEvent historyEvent in historyEvents) + { + sizeInBytes += historyEvent.CalculateSize(); + } + + return sizeInBytes; + } + + static bool IsOrchestrationFinished(OrchestrationStatus orchestrationStatus) + { + return orchestrationStatus is OrchestrationStatus.Completed + or OrchestrationStatus.Terminated + or OrchestrationStatus.Failed + or OrchestrationStatus.Canceled; } async ValueTask BuildRuntimeStateAsync( @@ -183,35 +222,53 @@ async ValueTask BuildRuntimeStateAsync( ? ProtoUtils.ConvertHistoryEvent : entityConversionState.ConvertFromProto; - IEnumerable pastEvents = []; - if (orchestratorRequest.RequiresHistoryStreaming) - { - // Stream the remaining events from the remote service - P.StreamInstanceHistoryRequest streamRequest = new() - { - InstanceId = orchestratorRequest.InstanceId, - ExecutionId = orchestratorRequest.ExecutionId, - ForWorkItemProcessing = true, - }; - - using AsyncServerStreamingCall streamResponse = - this.client.StreamInstanceHistory(streamRequest, cancellationToken: cancellation); - - await foreach (P.HistoryChunk chunk in streamResponse.ResponseStream.ReadAllAsync(cancellation)) - { - pastEvents = pastEvents.Concat(chunk.Events.Select(converter)); - } - } - else - { - // The history was already provided in the work item request - pastEvents = orchestratorRequest.PastEvents.Select(converter); - } - + IEnumerable? pastEvents = []; + + if (orchestratorRequest.PastEvents.Count > 0) + { + // The history was already provided in the work item request + pastEvents = orchestratorRequest.PastEvents.Select(converter); + int sizeInBytes = CalculateSizeInBytes(orchestratorRequest.PastEvents); + + // If the backend provided a history, even if we already have one cached, we should always store it + this.orchestrationHistories.Put(orchestratorRequest.InstanceId, pastEvents, sizeInBytes); + } + + // Either the request requires history streaming, or a history exists and none was provided in the request and nothing is cached so we have to stream it + else if (orchestratorRequest.RequiresHistoryStreaming || (orchestratorRequest.HistoryExists && !this.orchestrationHistories.TryGetValue(orchestratorRequest.InstanceId, out pastEvents))) + { + // Stream the remaining events from the remote service + P.StreamInstanceHistoryRequest streamRequest = new() + { + InstanceId = orchestratorRequest.InstanceId, + ExecutionId = orchestratorRequest.ExecutionId, + ForWorkItemProcessing = true, + }; + + using AsyncServerStreamingCall streamResponse = + this.client.StreamInstanceHistory(streamRequest, cancellationToken: cancellation); + + pastEvents = []; + int sizeInBytes = 0; + await foreach (P.HistoryChunk chunk in streamResponse.ResponseStream.ReadAllAsync(cancellation)) + { + pastEvents = pastEvents.Concat(chunk.Events.Select(converter)); + sizeInBytes += CalculateSizeInBytes(chunk.Events); + } + + this.orchestrationHistories.Put(orchestratorRequest.InstanceId, pastEvents, sizeInBytes); + } + else + { + // Even if there is no history, we still want to add an item to the cache. This is necessary because upon the completion of the work item, the new history + // will only be added if an entry already exists for this instance id in the cache. + this.orchestrationHistories.Put(orchestratorRequest.InstanceId, [], 0); + } + IEnumerable newEvents = orchestratorRequest.NewEvents.Select(converter); // Reconstruct the orchestration state in a way that correctly distinguishes new events from past events - var runtimeState = new OrchestrationRuntimeState(pastEvents.ToList()); + var runtimeState = new OrchestrationRuntimeState(pastEvents!.ToList()); foreach (HistoryEvent e in newEvents) { // AddEvent() puts events into the NewEvents list. @@ -232,7 +289,7 @@ async ValueTask BuildRuntimeStateAsync( await this.client!.HelloAsync(EmptyMessage, cancellationToken: cancellation); this.Logger.EstablishedWorkItemConnection(); - DurableTaskWorkerOptions workerOptions = this.worker.workerOptions; + DurableTaskWorkerOptions workerOptions = this.worker.workerOptions; // Get the stream for receiving work-items return this.client!.GetWorkItems( @@ -245,7 +302,7 @@ async ValueTask BuildRuntimeStateAsync( MaxConcurrentEntityWorkItems = workerOptions.Concurrency.MaximumConcurrentEntityWorkItems, Capabilities = { P.WorkerCapability.HistoryStreaming }, - }, + }, cancellationToken: cancellation); } @@ -282,7 +339,10 @@ async Task ProcessWorkItemsAsync(AsyncServerStreamingCall stream, Ca cancellation)); } else if (workItem.RequestCase == P.WorkItem.RequestOneofCase.EntityRequest) - { + { + // Note that in this case, we call run entity batch without setting whether or not the entity state exists, which is by default false. + // This is the desired behavior since this type of entity request does not have worker state caching implemented, and so if an entity state is not provided + // in the request, that means that no state exists. this.RunBackgroundTask( workItem, () => this.OnRunEntityBatchAsync(workItem.EntityRequest.ToEntityBatchRequest(), cancellation)); @@ -299,7 +359,8 @@ async Task ProcessWorkItemsAsync(AsyncServerStreamingCall stream, Ca batchRequest, cancellation, workItem.CompletionToken, - operationInfos)); + operationInfos, + workItem.EntityRequestV2.EntityStateExists)); } else if (workItem.RequestCase == P.WorkItem.RequestOneofCase.HealthPing) { @@ -427,7 +488,8 @@ async Task OnRunOrchestratorAsync( failureDetails = unexpected.ToTaskFailureDetails(); } - P.OrchestratorResponse response; + P.OrchestratorResponse response; + OrchestrationStatus orchestrationStatus; if (result != null) { response = ProtoUtils.ConstructOrchestratorResponse( @@ -435,13 +497,16 @@ async Task OnRunOrchestratorAsync( result.CustomStatus, result.Actions, completionToken, - entityConversionState); + entityConversionState, + out orchestrationStatus, + historyCached: this.orchestrationHistories.ContainsKey(request.InstanceId)); // It could be that the history was never cached if it was "too big" (see comment for "Put" method) } else if (versioning != null && failureDetails != null && versionFailure) { - this.Logger.OrchestrationVersionFailure(versioning.FailureStrategy.ToString(), failureDetails.ErrorMessage); + this.Logger.OrchestrationVersionFailure(versioning.FailureStrategy.ToString(), failureDetails.ErrorMessage); if (versioning.FailureStrategy == DurableTaskWorkerOptions.VersionFailureStrategy.Fail) - { + { + orchestrationStatus = OrchestrationStatus.Failed; response = new P.OrchestratorResponse { InstanceId = request.InstanceId, @@ -473,8 +538,9 @@ await this.client.AbandonTaskOrchestratorWorkItemAsync( } } else - { - // This is the case for failures that happened *outside* the orchestrator executor + { + // This is the case for failures that happened *outside* the orchestrator executor + orchestrationStatus = OrchestrationStatus.Failed; response = new P.OrchestratorResponse { InstanceId = request.InstanceId, @@ -497,9 +563,38 @@ await this.client.AbandonTaskOrchestratorWorkItemAsync( name, response.InstanceId, response.Actions.Count, - GetActionsListForLogging(response.Actions)); - - await this.client.CompleteOrchestratorTaskAsync(response, cancellationToken: cancellationToken); + GetActionsListForLogging(response.Actions)); + + Func converter = entityConversionState is null + ? ProtoUtils.ConvertHistoryEvent + : entityConversionState.ConvertFromProto; + + try + { + var completeOrchestratorTaskResponse = await this.client.CompleteOrchestratorTaskAsync(response, cancellationToken: cancellationToken); + if (!IsOrchestrationFinished(orchestrationStatus)) + { + // If the entry has since been evicted from the cache, we do not want to store the new history for it since it will potentially be incomplete if the entry had an existing history attached. + if (this.orchestrationHistories.TryGetValueWithSize(request.InstanceId, out (IEnumerable? Value, int Size) pastEventsWithSize)) + { + int sizeInBytes = CalculateSizeInBytes(completeOrchestratorTaskResponse.NewHistory) + pastEventsWithSize.Size; + this.orchestrationHistories.Put( + request.InstanceId, + pastEventsWithSize.Value!.Concat(completeOrchestratorTaskResponse.NewHistory.Select(converter)), + sizeInBytes); + } + } + else + { + this.orchestrationHistories.Remove(request.InstanceId); + } + } + catch (RpcException e) when (e.StatusCode == StatusCode.NotFound) + { + // The instance was not found - this can happen if another worker completed the item due to a network failure, for example + // In this case, we want to remove the orchestration history from the cache since we are no longer responsible for this instance ID + this.orchestrationHistories.Remove(request.InstanceId); + } } async Task OnRunActivityAsync(P.ActivityRequest request, string completionToken, CancellationToken cancellation) @@ -546,7 +641,7 @@ async Task OnRunActivityAsync(P.ActivityRequest request, string completionToken, } else if (output != null) { - outputSizeInBytes = Encoding.UTF8.GetByteCount(output); + outputSizeInBytes = output.Length * sizeof(char); } string successOrFailure = failureDetails != null ? "failure" : "success"; @@ -569,26 +664,68 @@ async Task OnRunEntityBatchAsync( EntityBatchRequest batchRequest, CancellationToken cancellation, string? completionToken = null, - List? operationInfos = null) + List? operationInfos = null, + bool entityStateExists = false) { var coreEntityId = DTCore.Entities.EntityId.FromString(batchRequest.InstanceId!); EntityId entityId = new(coreEntityId.Name, coreEntityId.Key); TaskName name = new(entityId.Name); - EntityBatchResult? batchResult; + EntityBatchResult? batchResult; + bool entityStateCached = false; try { await using AsyncServiceScope scope = this.worker.services.CreateAsyncScope(); - IDurableTaskFactory2 factory = (IDurableTaskFactory2)this.worker.Factory; + IDurableTaskFactory2 factory = (IDurableTaskFactory2)this.worker.Factory; if (factory.TryCreateEntity(name, scope.ServiceProvider, out ITaskEntity? entity)) { // Both the factory invocation and the RunAsync could involve user code and need to be handled as // part of try/catch. - TaskEntity shim = this.shimFactory.CreateEntity(name, entity, entityId); - batchResult = await shim.ExecuteOperationBatchAsync(batchRequest); + TaskEntity shim = this.shimFactory.CreateEntity(name, entity, entityId); + + string? entityState = null; + + // If the backend provided an entity state, even if we already have one cached, we should always use it + if (batchRequest.EntityState == null) + { + if (entityStateExists && !this.entityStates.TryGetValue(batchRequest.InstanceId!, out entityState)) + { + P.GetEntityResponse getEntityResponse = await this.client.GetEntityAsync( + new P.GetEntityRequest + { + InstanceId = batchRequest.InstanceId, + IncludeState = true, + }, + cancellationToken: cancellation); + + // It should never be possible that this is false if entityStateExists is true, but we check it just in case. + if (getEntityResponse.Exists) + { + batchRequest.EntityState = getEntityResponse.Entity.SerializedState; + } + } + else + { + // In this case, either we successfully extracted the state from the cache, or a state does not exist, meaning the field should be null. + batchRequest.EntityState = entityState; + } + } + + batchResult = await shim.ExecuteOperationBatchAsync(batchRequest); + + // Even if the entity state is the same, we want to issue a put to refresh its position in the cache. + if (batchResult.EntityState != null) + { + entityStateCached = this.entityStates.Put(batchRequest.InstanceId!, batchResult.EntityState, batchResult.EntityState.Length * sizeof(char)); + } + else + { + // If the entity state is now null, remove whatever state we had in the cache for it, if any. + this.entityStates.Remove(batchRequest.InstanceId!); + } } else { @@ -610,7 +747,8 @@ async Task OnRunEntityBatchAsync( }, batchRequest.Operations!.Count).ToList(), FailureDetails = null, - }; + }; + this.entityStates.Remove(batchRequest.InstanceId!); } } catch (Exception frameworkException) @@ -626,7 +764,8 @@ async Task OnRunEntityBatchAsync( P.EntityBatchResult response = batchResult.ToEntityBatchResult( completionToken, - operationInfos?.Take(batchResult.Results?.Count ?? 0)); + operationInfos?.Take(batchResult.Results?.Count ?? 0), + entityStateCached); await this.client.CompleteEntityTaskAsync(response, cancellationToken: cancellation); } diff --git a/src/Worker/Grpc/GrpcDurableTaskWorker.cs b/src/Worker/Grpc/GrpcDurableTaskWorker.cs index cb11054d..618cc1c7 100644 --- a/src/Worker/Grpc/GrpcDurableTaskWorker.cs +++ b/src/Worker/Grpc/GrpcDurableTaskWorker.cs @@ -48,7 +48,8 @@ protected override async Task ExecuteAsync(CancellationToken stoppingToken) { await using AsyncDisposable disposable = this.GetCallInvoker(out CallInvoker callInvoker, out string address); this.logger.StartingTaskHubWorker(address); - await new Processor(this, new(callInvoker)).ExecuteAsync(stoppingToken); + using var processor = new Processor(this, new(callInvoker)); + await processor.ExecuteAsync(stoppingToken); } #if NET6_0_OR_GREATER diff --git a/src/Worker/Grpc/GrpcOrchestrationRunner.cs b/src/Worker/Grpc/GrpcOrchestrationRunner.cs index 9a6e3c5f..e0bf2c6b 100644 --- a/src/Worker/Grpc/GrpcOrchestrationRunner.cs +++ b/src/Worker/Grpc/GrpcOrchestrationRunner.cs @@ -119,8 +119,9 @@ public static string LoadAndRun( request.InstanceId, result.CustomStatus, result.Actions, - completionToken: string.Empty, /* doesn't apply */ - entityConversionState: null); + completionToken: string.Empty, /* doesn't apply */ + entityConversionState: null, + out OrchestrationStatus orchestrationStatus); byte[] responseBytes = response.ToByteArray(); return Convert.ToBase64String(responseBytes); } diff --git a/src/Worker/Grpc/LRUCache.cs b/src/Worker/Grpc/LRUCache.cs new file mode 100644 index 00000000..0ad50274 --- /dev/null +++ b/src/Worker/Grpc/LRUCache.cs @@ -0,0 +1,222 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +namespace Microsoft.DurableTask.Worker.Grpc; + +/// +/// Represents a Least Recently Used (LRU) cache of key-value pairs with a fixed capacity and a stale eviction time. +/// If the cache is at capacity and a new item is added, the least recently used item will be evicted. +/// The cache will be periodicially checked for items that have not been accessed within some timeframe, and evict +/// these "stale" items. +/// +/// The type of the key. +/// The type of the value. +class LRUCache : IDisposable where TKey : notnull +{ + const int MaximumPercentageOfCacheSizePerItem = 75; // Maximum percentage of the cache that a single item can take up. + + readonly int capacityInBytes; + readonly int checkForStaleItemsPeriodInMilliseconds; + readonly int staleEvectionTimeInMilliseconds; + readonly Dictionary> cacheMap = []; + readonly LinkedList<(TKey Key, TimedEntryWithSize TimedEntryWithSize)> lruList = new(); + readonly object cacheLock = new(); + readonly Timer staleEvictionTimer; + + int sizeInBytes; + + /// + /// Initializes a new instance of the class with the specified capacity and stale eviction time. + /// + /// The size of the cache, in bytes. + /// The amount of time between the evictions of stale items. + /// The amount of time that key-value pairs remain in the cache before they are considered "stale" + /// and evicted in the next check for stale items - a call to Put or TryGet on the key will reset the time. + /// Thrown if any of the parameters are less than or equal to 0. + internal LRUCache(int capacityInBytes, int checkForStaleItemsPeriodInMilliseconds, int staleEvectionTimeInMilliseconds) + { + if (capacityInBytes <= 0) + { + throw new ArgumentOutOfRangeException(nameof(capacityInBytes), "Capacity of the LRU cache must be greater than 0 bytes"); + } + + if (checkForStaleItemsPeriodInMilliseconds <= 0) + { + throw new ArgumentOutOfRangeException(nameof(checkForStaleItemsPeriodInMilliseconds), "The amount of time between checks for stale items must be greater than 0 milliseconds"); + } + + if (staleEvectionTimeInMilliseconds <= 0) + { + throw new ArgumentOutOfRangeException(nameof(staleEvectionTimeInMilliseconds), "The amount of time before an item is considered stale and eligible for eviction must be greater than 0 milliseconds"); + } + + this.capacityInBytes = capacityInBytes; + this.sizeInBytes = 0; + this.checkForStaleItemsPeriodInMilliseconds = checkForStaleItemsPeriodInMilliseconds; + this.staleEvectionTimeInMilliseconds = staleEvectionTimeInMilliseconds; + this.staleEvictionTimer = new(this.EvictStaleItems, null, this.staleEvectionTimeInMilliseconds, this.checkForStaleItemsPeriodInMilliseconds); + } + + /// + /// Disposes the LRU cache. + /// + public void Dispose() + { + this.staleEvictionTimer.Dispose(); + } + + /// + /// Returns the number of bytes the cache is able to accommodate, calculated by subtracting the current size in bytes from the capacity in bytes. + /// + /// The number of bytes remaining free in the cache. + internal int GetFreeSpaceInBytes() + { + lock (this.cacheLock) + { + return this.capacityInBytes - this.sizeInBytes; + } + } + + /// + /// Checks if the cache contains this key. + /// + /// The key to check. + /// True if the cache contains the key; otherwise false. + internal bool ContainsKey(TKey key) + { + lock (this.cacheLock) + { + return this.cacheMap.ContainsKey(key); + } + } + + /// + /// Attempts to retrieve the value associated with the specified key from the cache. If the key is in the cache, its value is assigned to the value parameter + /// and the key is moved to the most recently used position in the cache and its eviction time is reset. If the key is not found, the default value for the type + /// of the value parameter is assigned to the value parameter instead. + /// + /// The key of the value to retrieve. + /// When this method returns, contains the value associated with the specified key and its size, if the key is found; + /// otherwise, the default value for the type of the value parameter. This parameter is passed uninitialized. + /// True if the cache contains an entry with the specified key; otherwise, false. + internal bool TryGetValueWithSize(TKey key, out (TValue? Value, int Size) valueWithSize) + { + lock (this.cacheLock) + { + if (this.cacheMap.TryGetValue(key, out LinkedListNode<(TKey Key, TimedEntryWithSize TimedEntryWithSize)>? node)) + { + this.lruList.Remove(node); + this.lruList.AddFirst(node); + node.Value.TimedEntryWithSize.LastAccessed = DateTimeOffset.UtcNow; + valueWithSize = (node.Value.TimedEntryWithSize.Value, node.Value.TimedEntryWithSize.SizeInBytes); + return true; + } + + valueWithSize = default; + return false; + } + } + + /// + /// Attempts to retrieve the value associated with the specified key from the cache. If the key is in the cache, its value is assigned to the value parameter + /// and the key is moved to the most recently used position in the cache and its eviction time is reset. If the key is not found, the default value for the type + /// of the value parameter is assigned to the value parameter instead. + /// + /// The key of the value to retrieve. + /// When this method returns, contains the value associated with the specified key, if the key is found; + /// otherwise, the default value for the type of the value parameter. This parameter is passed uninitialized. + /// True if the cache contains an entry with the specified key; otherwise, false. + internal bool TryGetValue(TKey key, out TValue? value) + { + bool keyFound = this.TryGetValueWithSize(key, out (TValue? Value, int Size) valueWithSize); + value = valueWithSize.Value; + return keyFound; + } + + /// + /// Adds a new key-value pair to the cache or updates the value of an existing key. If adding the item would exceed the cache's capacity, the least recently used items are + /// removed until there is enough space for the new item. The key will be moved to the most recently used position in the cache and its eviction time reset. + /// + /// The key to add. + /// The value associated with the key to add. + /// The size of the item in bytes. If the item is larger than the capacity of the cache, it will not be added to the cache. + /// True if the item was successfully stored, false if it was too big (>= than of the cache size ). + internal bool Put(TKey key, TValue value, int itemSizeInBytes) + { + lock (this.cacheLock) + { + // If the item is too big, we do not store it to avoid thrashing + if (itemSizeInBytes >= MaximumPercentageOfCacheSizePerItem / 100 * this.capacityInBytes) + { + return false; + } + + // If the key already exists, we update the value and move it to the front of the list + if (this.cacheMap.TryGetValue(key, out LinkedListNode<(TKey Key, TimedEntryWithSize TimedEntry)>? nodeToRemove)) + { + this.lruList.Remove(nodeToRemove); + } + + // If we will exceed capacity, keep removing the last nodes in the list (the least recently used items) until there is enough space. + // If the item's size in bytes is larger than the capacity, we will not add it to the cache. + else if (itemSizeInBytes <= this.capacityInBytes && this.sizeInBytes + itemSizeInBytes > this.capacityInBytes) + { + while (this.lruList.Last != null && this.sizeInBytes + itemSizeInBytes > this.capacityInBytes) + { + LinkedListNode<(TKey Key, TimedEntryWithSize TimedEntryWithSize)> lastNode = this.lruList.Last; + this.cacheMap.Remove(lastNode.Value.Key); + this.sizeInBytes -= lastNode.Value.TimedEntryWithSize.SizeInBytes; + this.lruList.RemoveLast(); + } + } + + var newNode = new LinkedListNode<(TKey Key, TimedEntryWithSize TimedEntry)>((key, new TimedEntryWithSize(value, DateTimeOffset.UtcNow, itemSizeInBytes))); + this.sizeInBytes += itemSizeInBytes; + this.lruList.AddFirst(newNode); + this.cacheMap[key] = newNode; + return true; + } + } + + /// + /// Removes the key and its value from the cache. If the key is not found, this method does nothing. + /// + /// The key to remove from the cache. + internal void Remove(TKey key) + { + lock (this.cacheLock) + { + if (this.cacheMap.TryGetValue(key, out LinkedListNode<(TKey Key, TimedEntryWithSize TimedEntry)>? nodeToRemove)) + { + this.sizeInBytes -= nodeToRemove.Value.TimedEntry.SizeInBytes; + this.cacheMap.Remove(key); + this.lruList.Remove(nodeToRemove); + } + } + } + + void EvictStaleItems(object? state) + { + var now = DateTimeOffset.UtcNow; + lock (this.cacheLock) + { + // Since the list is ordered from most recently used to least recently used, we can iterate from the end of the list until we reach an item that is not stale. + // All items preceding it are guaranteed to not be stale as well. + while (this.lruList.Last != null && (now - this.lruList.Last.Value.TimedEntryWithSize.LastAccessed).TotalMilliseconds >= this.staleEvectionTimeInMilliseconds) + { + this.sizeInBytes -= this.lruList.Last.Value.TimedEntryWithSize.SizeInBytes; + this.cacheMap.Remove(this.lruList.Last.Value.Key); + this.lruList.RemoveLast(); + } + } + } + + class TimedEntryWithSize(TValue value, DateTimeOffset lastAccessed, int sizeInBytes) + { + public TValue Value { get; set; } = value; + + public DateTimeOffset LastAccessed { get; set; } = lastAccessed; + + public int SizeInBytes { get; set; } = sizeInBytes; + } +}