ArchiveDocumentManager.cs

namespace ServiceControl.Recoverability
{
    using System;
    using System.Collections.Generic;
    using System.Linq;
    using System.Threading;
    using System.Threading.Tasks;
    using MessageFailures;
    using Microsoft.Extensions.Logging;
    using Persistence.RavenDB;
    using Raven.Client.Documents;
    using Raven.Client.Documents.Commands.Batches;
    using Raven.Client.Documents.Operations;
    using Raven.Client.Documents.Session;

    class ArchiveDocumentManager(ExpirationManager expirationManager, ILogger logger, TimeProvider timeProvider)
    {
        public Task<ArchiveOperation> LoadArchiveOperation(IAsyncDocumentSession session, string groupId, ArchiveType archiveType, CancellationToken cancellationToken = default) => session.LoadAsync<ArchiveOperation>(ArchiveOperation.MakeId(groupId, archiveType), cancellationToken);

        public async Task<ArchiveOperation> CreateArchiveOperation(IAsyncDocumentSession session, string groupId, ArchiveType archiveType, int numberOfMessages, string groupName, int batchSize, string initiatedById = null, string initiatedByName = null, string operationId = null, CancellationToken cancellationToken = default)
        {
            var operation = new ArchiveOperation
            {
                Id = ArchiveOperation.MakeId(groupId, archiveType),
                RequestId = groupId,
                ArchiveType = archiveType,
                TotalNumberOfMessages = numberOfMessages,
                NumberOfMessagesArchived = 0,
                Started = timeProvider.GetUtcNow().UtcDateTime,
                GroupName = groupName,
                NumberOfBatches = (int)Math.Ceiling(numberOfMessages / (float)batchSize),
                CurrentBatch = 0,
                InitiatedById = initiatedById,
                InitiatedByName = initiatedByName,
                OperationId = operationId
            };

            await session.StoreAsync(operation, cancellationToken);

            var documentCount = 0;
            var indexQuery = session.Query<FailureGroupMessageView>(new FailedMessages_ByGroup().IndexName);

            var docQuery = indexQuery
                .Where(failure => failure.FailureGroupId == groupId)
                .Where(failure => failure.Status == FailedMessageStatus.Unresolved)
                .Select(document => document.Id);

            var docs = await StreamResults(session, docQuery, cancellationToken);

            var batches = docs
                .GroupBy(d => documentCount++ / batchSize);

            foreach (var batch in batches)
            {
                var archiveBatch = new ArchiveBatch
                {
                    Id = ArchiveBatch.MakeId(groupId, archiveType, batch.Key),
                    DocumentIds = batch.ToList()
                };

                await session.StoreAsync(archiveBatch, cancellationToken);
            }

            return operation;
        }

        async Task<IEnumerable<string>> StreamResults(IAsyncDocumentSession session, IQueryable<string> query, CancellationToken cancellationToken)
        {
            var results = new List<string>();
            await using var enumerator = await session.Advanced.StreamAsync(query, cancellationToken);
            while (await enumerator.MoveNextAsync())
            {
                results.Add(enumerator.Current.Document);
            }

            return results;
        }

        public Task<ArchiveBatch> GetArchiveBatch(IAsyncDocumentSession session, string archiveOperationId, int batchNumber, CancellationToken cancellationToken = default) => session.LoadAsync<ArchiveBatch>($"{archiveOperationId}/{batchNumber}", cancellationToken);

        public async Task<GroupDetails> GetGroupDetails(IAsyncDocumentSession session, string groupId, CancellationToken cancellationToken = default)
        {
            var group = await session.Query<FailureGroupView, FailureGroupsViewIndex>()
                .FirstOrDefaultAsync(x => x.Id == groupId, cancellationToken);

            return new GroupDetails
            {
                NumberOfMessagesInGroup = group?.Count ?? 0,
                GroupName = group?.Title ?? "Undefined"
            };
        }

        public void ArchiveMessageGroupBatch(IAsyncDocumentSession session, ArchiveBatch batch)
        {
            var patchRequest = new PatchRequest
            {
                Script = "this.Status = args.Status;",
                Values =
                {
                    { "Status", (int)FailedMessageStatus.Archived }
                }
            };

            expirationManager.EnableExpiration(patchRequest);

            var patchCommands = batch?.DocumentIds.Select(documentId => new PatchCommandData(documentId, null, patchRequest));

            if (patchCommands != null)
            {
                session.Advanced.Defer(patchCommands.ToArray<ICommandData>());
                session.Advanced.Defer(new DeleteCommandData(batch.Id, null));
            }
        }

        public async Task<bool> WaitForIndexUpdateOfArchiveOperation(IRavenSessionProvider sessionProvider, string requestId, TimeSpan timeToWait, CancellationToken cancellationToken = default)
        {
            using var session = await sessionProvider.OpenSession(cancellationToken: cancellationToken);
            var indexQuery = session.Query<FailureGroupMessageView>(new FailedMessages_ByGroup().IndexName)
                .Customize(x => x.WaitForNonStaleResults(timeToWait));

            var docQuery = indexQuery
                .Where(failure => failure.FailureGroupId == requestId)
                .Select(document => document.Id);

            try
            {
                await docQuery.AnyAsync(cancellationToken);

                return true;
            }
            catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
            {
                throw;
            }
            catch
            {
                return false;
            }
        }

        public Task UpdateArchiveOperation(IAsyncDocumentSession session, ArchiveOperation archiveOperation, CancellationToken cancellationToken = default) => session.StoreAsync(archiveOperation, cancellationToken);

        public async Task RemoveArchiveOperation(IRavenSessionProvider sessionProvider, ArchiveOperation archiveOperation, CancellationToken cancellationToken = default)
        {
            using var session = await sessionProvider.OpenSession(cancellationToken: cancellationToken);
            session.Advanced.Defer(new DeleteCommandData(archiveOperation.Id, null));
            await session.SaveChangesAsync(cancellationToken);

            logger.LogInformation("Removing ArchiveOperation {ArchiveOperationId} completed", archiveOperation.Id);
        }

        public class GroupDetails
        {
            public string GroupName { get; set; }
            public int NumberOfMessagesInGroup { get; set; }
        }

    }
}