fix: merging dev to main

App/backend-api/Microsoft.GS.DPS/API/KernelMemory/KernelMemory.cs

@@ -1,16 +1,19 @@
 using Microsoft.Extensions.Logging;
 using Microsoft.GS.DPS.Images;
 using Microsoft.GS.DPS.Model.KernelMemory;
+using Microsoft.GS.DPS.Storage.Components;
 using Microsoft.GS.DPS.Storage.Document;
 using Microsoft.KernelMemory;
 using Microsoft.KernelMemory.Context;
 using Microsoft.KernelMemory.Pipeline;
 using MongoDB.Bson;
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Data;
 using System.Linq;
 using System.Reflection;
+using System.Security.Cryptography;
 using System.Text;
 using System.Text.Json;
 using System.Threading.Tasks;
@@ -28,7 +31,11 @@ public class KernelMemory
         private readonly DataCacheManager _dataCache;
         private readonly TagUpdater _tagUpdator;
         private readonly ILogger<KernelMemory>? _logger;
+        private readonly ConcurrentDictionary<string, Lazy<Task<DocumentImportedResult>>> _documentImports = new();
         private static readonly string keywordExtractorPrompt = "";
+        private static readonly TimeSpan importLeaseDuration = TimeSpan.FromMinutes(10);
+        private static readonly TimeSpan importLeaseRenewalInterval = TimeSpan.FromMinutes(1);
+        private static readonly TimeSpan importLeaseWaitTimeout = TimeSpan.FromMinutes(70);
 
         static KernelMemory()
         {
@@ -53,8 +60,94 @@ public async Task<DocumentImportedResult> ImportDocument(Stream documentStream,
                                                                  string fileName, 
                                                                  string contentType)
         {
-            // Implementation of the file upload
-            var documentId = await _kmClient.ImportDocumentAsync(documentStream, fileName, steps: [
+            using var bufferedStream = documentStream.CanSeek ? null : CreateTemporaryFileStream();
+            Stream importStream = documentStream;
+
+            if (bufferedStream != null)
+            {
+                await documentStream.CopyToAsync(bufferedStream);
+                importStream = bufferedStream;
+            }
+
+            importStream.Position = 0;
+            var contentHash = await SHA256.HashDataAsync(importStream);
+            importStream.Position = 0;
+            var documentId = Convert.ToHexString(contentHash).ToLowerInvariant();
+
+            var documentImport = _documentImports.GetOrAdd(
+                documentId,
+                _ => new Lazy<Task<DocumentImportedResult>>(
+                    () => ImportDocumentCore(importStream, fileName, contentType, documentId),
+                    LazyThreadSafetyMode.ExecutionAndPublication));
+
+            try
+            {
+                return await documentImport.Value;
+            }
+            finally
+            {
+                ((ICollection<KeyValuePair<string, Lazy<Task<DocumentImportedResult>>>>)_documentImports)
+                    .Remove(new KeyValuePair<string, Lazy<Task<DocumentImportedResult>>>(documentId, documentImport));
+            }
+        }
+
+        private static FileStream CreateTemporaryFileStream()
+        {
+            var temporaryFilePath = Path.Combine(Path.GetTempPath(), Path.GetRandomFileName());
+            return new FileStream(
+                temporaryFilePath,
+                FileMode.CreateNew,
+                FileAccess.ReadWrite,
+                FileShare.None,
+                bufferSize: 81920,
+                FileOptions.Asynchronous | FileOptions.SequentialScan | FileOptions.DeleteOnClose);
+        }
+
+        private async Task<DocumentImportedResult> ImportDocumentCore(Stream importStream,
+                                                                       string fileName,
+                                                                       string contentType,
+                                                                       string documentId)
+        {
+            var existingDocument = await _documentRepository.FindByDocumentIdAsync(documentId);
+            if (existingDocument != null)
+            {
+                return ToImportedResult(existingDocument);
+            }
+
+            var leaseOwnerId = Guid.NewGuid().ToString("N");
+            var leaseWaitDeadline = DateTime.UtcNow.Add(importLeaseWaitTimeout);
+
+            while (!await _documentRepository.TryAcquireImportLeaseAsync(
+                       documentId,
+                       leaseOwnerId,
+                       DateTime.UtcNow.Add(importLeaseDuration)))
+            {
+                if (DateTime.UtcNow >= leaseWaitDeadline)
+                {
+                    throw new TimeoutException("Timed out waiting for another import of this document to complete.");
+                }
+
+                await Task.Delay(TimeSpan.FromSeconds(2));
+                existingDocument = await _documentRepository.FindByDocumentIdAsync(documentId);
+                if (existingDocument != null)
+                {
+                    return ToImportedResult(existingDocument);
+                }
+            }
+
+            using var leaseRenewalCancellation = new CancellationTokenSource();
+            var leaseRenewalTask = RenewImportLeaseAsync(documentId, leaseOwnerId, leaseRenewalCancellation.Token);
+
+            try
+            {
+                existingDocument = await _documentRepository.FindByDocumentIdAsync(documentId);
+                if (existingDocument != null)
+                {
+                    return ToImportedResult(existingDocument);
+                }
+
+                // Implementation of the file upload
+                await _kmClient.ImportDocumentAsync(importStream, fileName, documentId: documentId, steps: [
                                     Constants.PipelineStepsExtract,
                                     "keyword_extract",
                                     Constants.PipelineStepsSummarize,
@@ -97,6 +190,7 @@ public async Task<DocumentImportedResult> ImportDocument(Stream documentStream,
             // Save the document to the repository
             Document document = new Document
             {
+                id = new Guid(Convert.FromHexString(documentId[..32])),
                 DocumentId = documentId,
                 FileName = fileName,
                 ImportedTime = importedResult.ImportedTime,
@@ -105,13 +199,71 @@ public async Task<DocumentImportedResult> ImportDocument(Stream documentStream,
                 Summary = importedResult.Summary,
                 Keywords = importedResult.Keywords
             };
+            document.__partitionkey = CosmosDBEntityBase.GetKey(document.id, 9999);
+
+                await _documentRepository.RegisterAsync(document);
+
+                //Cache Refresh
+                _dataCache.ManualRefresh();
+
+                return importedResult;
+            }
+            finally
+            {
+                leaseRenewalCancellation.Cancel();
+
+                try
+                {
+                    await leaseRenewalTask;
+                }
+                catch (OperationCanceledException) when (leaseRenewalCancellation.IsCancellationRequested)
+                {
+                }
+                catch (Exception exception)
+                {
+                    _logger?.LogWarning(exception, "Failed to renew the import lease for document {DocumentId}", documentId);
+                }
+
+                try
+                {
+                    await _documentRepository.ReleaseImportLeaseAsync(documentId, leaseOwnerId);
+                }
+                catch (Exception exception)
+                {
+                    _logger?.LogWarning(exception, "Failed to release the import lease for document {DocumentId}", documentId);
+                }
+            }
+        }
 
-            await _documentRepository.RegisterAsync(document);
+        private async Task RenewImportLeaseAsync(string documentId, string ownerId, CancellationToken cancellationToken)
+        {
+            while (true)
+            {
+                await Task.Delay(importLeaseRenewalInterval, cancellationToken);
+                var renewed = await _documentRepository.RenewImportLeaseAsync(
+                    documentId,
+                    ownerId,
+                    DateTime.UtcNow.Add(importLeaseDuration));
 
-            //Cache Refresh
-            _dataCache.ManualRefresh();
+                if (!renewed)
+                {
+                    throw new InvalidOperationException($"The import lease for document {documentId} is no longer owned by this process.");
+                }
+            }
+        }
 
-            return importedResult;
+        private static DocumentImportedResult ToImportedResult(Document document)
+        {
+            return new DocumentImportedResult
+            {
+                DocumentId = document.DocumentId,
+                ImportedTime = document.ImportedTime,
+                MimeType = document.MimeType,
+                FileName = document.FileName,
+                ProcessingTime = document.ProcessingTime,
+                Keywords = document.Keywords,
+                Summary = document.Summary
+            };
         }
 
         public async Task<bool> DeleteDocument(string documentId)

App/backend-api/Microsoft.GS.DPS/Storage/Documents/DocumentRepository.cs

@@ -6,6 +6,7 @@
 using MongoDB.Driver;
 using System.ComponentModel;
 using MongoDB.Bson;
+using MongoDB.Bson.Serialization.Attributes;
 
 namespace Microsoft.GS.DPS.Storage.Document
 {
@@ -14,9 +15,20 @@ namespace Microsoft.GS.DPS.Storage.Document
     public class DocumentRepository 
     {
         private readonly IMongoCollection<Entities.Document> _collection;
+        private readonly IMongoCollection<DocumentImportLease> _importLeases;
+
+        private sealed class DocumentImportLease
+        {
+            [BsonId]
+            public string DocumentId { get; set; } = string.Empty;
+            public string OwnerId { get; set; } = string.Empty;
+            public DateTime ExpiresAt { get; set; }
+        }
+
         public DocumentRepository(IMongoDatabase database, string collectionName) 
         {
             _collection = database.GetCollection<Entities.Document>(collectionName);
+            _importLeases = database.GetCollection<DocumentImportLease>($"{collectionName}_ImportLeases");
 
             // if Database is empty, create a new collection
             if (_collection == null)
@@ -139,8 +151,85 @@ private int GetTotalPages(int pageSize, double recordsCount)
 
         public async Task<Entities.Document> RegisterAsync(Entities.Document document)
         {
-            await _collection.InsertOneAsync(document);
-            return document;
+            var existingDocument = await FindByDocumentIdAsync(document.DocumentId);
+            if (existingDocument != null)
+            {
+                document.id = existingDocument.id;
+                document.__partitionkey = existingDocument.__partitionkey;
+            }
+
+            var filter = Builders<Entities.Document>.Filter.Eq(x => x.id, document.id);
+            var update = Builders<Entities.Document>.Update
+                .Set(x => x.FileName, document.FileName)
+                .Set(x => x.ImportedTime, document.ImportedTime)
+                .Set(x => x.MimeType, document.MimeType)
+                .Set(x => x.ProcessingTime, document.ProcessingTime)
+                .Set(x => x.Summary, document.Summary)
+                .Set(x => x.Keywords, document.Keywords)
+                .SetOnInsert(x => x.DocumentId, document.DocumentId)
+                .SetOnInsert(x => x.__partitionkey, document.__partitionkey);
+
+            return await _collection.FindOneAndUpdateAsync(
+                filter,
+                update,
+                new FindOneAndUpdateOptions<Entities.Document>
+                {
+                    IsUpsert = true,
+                    ReturnDocument = ReturnDocument.After
+                });
+        }
+
+        public async Task<bool> TryAcquireImportLeaseAsync(string documentId, string ownerId, DateTime expiresAt)
+        {
+            var lease = new DocumentImportLease
+            {
+                DocumentId = documentId,
+                OwnerId = ownerId,
+                ExpiresAt = expiresAt
+            };
+
+            try
+            {
+                await _importLeases.InsertOneAsync(lease);
+                return true;
+            }
+            catch (MongoWriteException exception) when (exception.WriteError?.Category == ServerErrorCategory.DuplicateKey)
+            {
+            }
+            catch (MongoCommandException exception) when (exception.Code == 11000)
+            {
+            }
+
+            var expiredLeaseFilter = Builders<DocumentImportLease>.Filter.Eq(x => x.DocumentId, documentId) &
+                                     Builders<DocumentImportLease>.Filter.Lte(x => x.ExpiresAt, DateTime.UtcNow);
+            var update = Builders<DocumentImportLease>.Update
+                .Set(x => x.OwnerId, ownerId)
+                .Set(x => x.ExpiresAt, expiresAt);
+            var acquiredLease = await _importLeases.FindOneAndUpdateAsync(
+                expiredLeaseFilter,
+                update,
+                new FindOneAndUpdateOptions<DocumentImportLease>
+                {
+                    ReturnDocument = ReturnDocument.After
+                });
+
+            return acquiredLease?.OwnerId == ownerId;
+        }
+
+        public async Task ReleaseImportLeaseAsync(string documentId, string ownerId)
+        {
+            var filter = Builders<DocumentImportLease>.Filter.Eq(x => x.DocumentId, documentId) &
+                         Builders<DocumentImportLease>.Filter.Eq(x => x.OwnerId, ownerId);
+            await _importLeases.DeleteOneAsync(filter);
+        }
+
+        public async Task<bool> RenewImportLeaseAsync(string documentId, string ownerId, DateTime expiresAt)
+        {
+            var filter = Builders<DocumentImportLease>.Filter.Eq(x => x.DocumentId, documentId) &
+                         Builders<DocumentImportLease>.Filter.Eq(x => x.OwnerId, ownerId);
+            var update = Builders<DocumentImportLease>.Update.Set(x => x.ExpiresAt, expiresAt);
+            var result = await _importLeases.UpdateOneAsync(filter, update);
+            return result.MatchedCount == 1;
         }
 
         public async Task<Entities.Document> UpdateAsync(Entities.Document document)