Azure · kirankumarkolli · Nov 4, 2025 · Sep 12, 2025 · Sep 12, 2025 · Sep 15, 2025
@@ -5,32 +5,127 @@
 namespace Microsoft.Azure.Cosmos.Encryption.Custom
 {
     using System;
+    using System.Collections.Generic;
+    using System.Diagnostics;
+    using System.Threading;
 
     /// <summary>
-    /// This is an empty implementation of CosmosDiagnosticsContext which has been plumbed through the DataEncryptionKeyProvider and EncryptionContainer.
-    /// This may help adding diagnostics more easily in future.
+    /// Lightweight diagnostics context for Custom Encryption extension.
+    /// Records scope names, start/stop timestamps and exposes them for tests or future wiring into SDK diagnostics.
+    /// Uses <see cref="ActivitySource"/> so downstream telemetry (OpenTelemetry) can optionally subscribe.
     /// </summary>
     internal class CosmosDiagnosticsContext
     {
-        private static readonly CosmosDiagnosticsContext UnusedSingleton = new ();
-        private static readonly IDisposable UnusedScopeSingleton = new Scope();
+        private static readonly ActivitySource ActivitySource = new ("Microsoft.Azure.Cosmos.Encryption.Custom");
 
+        private readonly List<ScopeRecord> records = new (4);
+
+        internal CosmosDiagnosticsContext()
+        {
+        }
+
+        /// <summary>
+        /// Factory. A new instance is created per high-level operation to avoid cross-talk.
+        /// </summary>
         public static CosmosDiagnosticsContext Create(RequestOptions options)
         {
-            _ = options;
-            return CosmosDiagnosticsContext.UnusedSingleton;
+            _ = options; // Reserved for future correlation if RequestOptions ever carries a diagnostics handle.
+            return new CosmosDiagnosticsContext();
+        }
+
+        /// <summary>
+        /// Recorded scope metadata (immutable snapshot on scope dispose).
+        /// </summary>
+        internal readonly struct ScopeRecord
+        {
+            public ScopeRecord(string name, long startTimestamp, long elapsedTicks)
+            {
+                this.Name = name;
+                this.StartTimestamp = startTimestamp;
+                this.ElapsedTicks = elapsedTicks;
+            }
+
+            public string Name { get; }
+
+            public long StartTimestamp { get; }
+
+            public long ElapsedTicks { get; } // Stopwatch ticks
+
+            public TimeSpan Elapsed => TimeSpan.FromTicks(this.ElapsedTicks);
+        }
+
+        /// <summary>
+        /// Gets recorded scope names primarily for unit tests (copy snapshot each access).
+        /// </summary>
+        internal IReadOnlyList<string> Scopes
+        {
+            get
+            {
+                if (this.records.Count == 0)
+                {
+                    return Array.Empty<string>();
+                }
+
+                string[] names = new string[this.records.Count];
+                for (int i = 0; i < this.records.Count; i++)
+                {
+                    names[i] = this.records[i].Name;
+                }
+
+                return names;
+            }
+        }
+
+        public Scope CreateScope(string scope)
+        {
+            if (string.IsNullOrEmpty(scope))
+            {
+                return Scope.Noop; // returns default(struct) => no-op
+            }
+
+            // Only create Activity if there are listeners to avoid unnecessary allocations.
+            Activity activity = ActivitySource.HasListeners() ? ActivitySource.StartActivity(scope, ActivityKind.Internal) : null;
+            long startTicks = Stopwatch.GetTimestamp();
+            return new Scope(this, scope, startTicks, activity);
         }
 
-        public IDisposable CreateScope(string scope)
+        private void Record(string name, long startTicks, long elapsedTicks)
         {
-            _ = scope;
-            return CosmosDiagnosticsContext.UnusedScopeSingleton;
+            lock (this.records)
+            {
+                this.records.Add(new ScopeRecord(name, startTicks, elapsedTicks));
+            }
         }
 
-        private class Scope : IDisposable
+        public readonly struct Scope : IDisposable
         {
+            private readonly CosmosDiagnosticsContext owner;
+            private readonly string name;
+            private readonly long startTicks;
+            private readonly Activity activity;
+            private readonly bool enabled;
+
+            internal Scope(CosmosDiagnosticsContext owner, string name, long startTicks, Activity activity)
+            {
+                this.owner = owner;
+                this.name = name;
+                this.startTicks = startTicks;
+                this.activity = activity;
+                this.enabled = owner != null; // default struct (Noop) => owner null
+            }
+
+            internal static Scope Noop => default;
+
             public void Dispose()
             {
+                if (!this.enabled)
+                {
+                    return;
+                }
+
+                long elapsedTicks = Stopwatch.GetTimestamp() - this.startTicks;
+                this.owner.Record(this.name, this.startTicks, elapsedTicks);
+                this.activity?.Dispose();
             }
         }
     }

@@ -40,14 +40,7 @@ internal CosmosJsonDotNetSerializer(JsonSerializerSettings jsonSerializerSetting
         /// <returns>The object representing the deserialized stream</returns>
         public T FromStream<T>(Stream stream)
         {
-#if NET8_0_OR_GREATER
-            ArgumentNullException.ThrowIfNull(stream);
-#else
-            if (stream == null)
-            {
-                throw new ArgumentNullException(nameof(stream));
-            }
-#endif
+            ArgumentValidation.ThrowIfNull(stream);
 
             if (typeof(Stream).IsAssignableFrom(typeof(T)))
             {
@@ -87,6 +80,44 @@ public MemoryStream ToStream<T>(T input)
             return streamPayload;
         }
 
+        /// <summary>
+        /// Serializes an object directly into the provided output stream (which remains open) and rewinds it if seekable.
+        /// </summary>
+        /// <typeparam name="T">Type of object being serialized.</typeparam>
+        /// <param name="input">Object to serialize.</param>
+        /// <param name="output">Destination stream. Must be writable. The stream is not disposed by this method.</param>
+        /// <exception cref="ArgumentNullException">Thrown when <paramref name="output"/> is <c>null</c>.</exception>
+        /// <exception cref="ArgumentException">Thrown when <paramref name="output"/> is not writable.</exception>
+        /// <remarks>
+        /// This method serializes the object directly to the provided stream without creating an intermediate MemoryStream,
+        /// reducing memory allocations for large objects. If the output stream is seekable, its position is reset to 0 after writing.
+        /// </remarks>
+        public void WriteToStream<T>(T input, Stream output)
+        {
+            ArgumentValidation.ThrowIfNull(output);
+
+            if (!output.CanWrite)
+            {
+                throw new ArgumentException("Output stream must be writable", nameof(output));
+            }
+
+            using (StreamWriter streamWriter = new (output, encoding: CosmosJsonDotNetSerializer.DefaultEncoding, bufferSize: 1024, leaveOpen: true))
+            using (JsonTextWriter writer = new (streamWriter))
+            {
+                writer.ArrayPool = JsonArrayPool.Instance;
+                writer.Formatting = Newtonsoft.Json.Formatting.None;
+                JsonSerializer jsonSerializer = this.GetSerializer();
+                jsonSerializer.Serialize(writer, input);
+                writer.Flush();
+                streamWriter.Flush();
+            }
+
+            if (output.CanSeek)
+            {
+                output.Position = 0;
+            }
+        }
+
         /// <summary>
         /// JsonSerializer has hit a race conditions with custom settings that cause null reference exception.
         /// To avoid the race condition a new JsonSerializer is created for each call

@@ -0,0 +1,31 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Encryption.Custom
+{
+    using System.IO;
+    using System.Threading.Tasks;
+
+    /// <summary>
+    /// Extension methods for Stream to provide compatibility across different .NET versions.
+    /// </summary>
+    internal static class StreamExtensions
+    {
+        /// <summary>
+        /// Asynchronously disposes the stream in a version-compatible way.
+        /// Uses DisposeAsync on .NET 8.0+ and falls back to synchronous Dispose on earlier versions.
+        /// </summary>
+        /// <param name="stream">The stream to dispose.</param>
+        /// <returns>A ValueTask representing the asynchronous dispose operation.</returns>
+        public static ValueTask DisposeCompatAsync(this Stream stream)
+        {
+#if NET8_0_OR_GREATER
+            return stream.DisposeAsync();
+#else
+            stream.Dispose();
+            return default;
+#endif
+        }
+    }
+}
@@ -148,14 +148,7 @@ public async Task InitializeAsync(
                 throw new InvalidOperationException($"{nameof(CosmosDataEncryptionKeyProvider)} has already been initialized.");
             }
 
-#if NET8_0_OR_GREATER
-            ArgumentNullException.ThrowIfNull(database);
-#else
-            if (database == null)
-            {
-                throw new ArgumentNullException(nameof(database));
-            }
-#endif
+            ArgumentValidation.ThrowIfNull(database);
 
             ContainerResponse containerResponse = await database.CreateContainerIfNotExistsAsync(
                 containerId,

@@ -102,14 +102,7 @@ public static DataEncryptionKey Create(
             byte[] rawKey,
             string encryptionAlgorithm)
         {
-#if NET8_0_OR_GREATER
-            ArgumentNullException.ThrowIfNull(rawKey);
-#else
-            if (rawKey == null)
-            {
-                throw new ArgumentNullException(nameof(rawKey));
-            }
-#endif
+            ArgumentValidation.ThrowIfNull(rawKey);
 
 #pragma warning disable CS0618 // Type or member is obsolete
             if (!string.Equals(encryptionAlgorithm, CosmosEncryptionAlgorithm.AEAes256CbcHmacSha256Randomized))

@@ -67,14 +67,7 @@ public override async Task<ItemResponse<DataEncryptionKeyProperties>> CreateData
                 throw new ArgumentException(string.Format("Unsupported Encryption Algorithm {0}", encryptionAlgorithm), nameof(encryptionAlgorithm));
             }
 
-#if NET8_0_OR_GREATER
-            ArgumentNullException.ThrowIfNull(encryptionKeyWrapMetadata);
-#else
-            if (encryptionKeyWrapMetadata == null)
-            {
-                throw new ArgumentNullException(nameof(encryptionKeyWrapMetadata));
-            }
-#endif
+            ArgumentValidation.ThrowIfNull(encryptionKeyWrapMetadata);
 
             CosmosDiagnosticsContext diagnosticsContext = CosmosDiagnosticsContext.Create(requestOptions);
 
@@ -159,14 +152,7 @@ public override async Task<ItemResponse<DataEncryptionKeyProperties>> RewrapData
            ItemRequestOptions requestOptions = null,
            CancellationToken cancellationToken = default)
         {
-#if NET8_0_OR_GREATER
-            ArgumentNullException.ThrowIfNull(newWrapMetadata);
-#else
-            if (newWrapMetadata == null)
-            {
-                throw new ArgumentNullException(nameof(newWrapMetadata));
-            }
-#endif
+            ArgumentValidation.ThrowIfNull(newWrapMetadata);
 
             CosmosDiagnosticsContext diagnosticsContext = CosmosDiagnosticsContext.Create(requestOptions);
 

@@ -78,14 +78,7 @@ public override Task<ItemResponse<DataEncryptionKeyProperties>> RewrapDataEncryp
                 throw new ArgumentNullException(nameof(id));
             }
 
-#if NET8_0_OR_GREATER
-            ArgumentNullException.ThrowIfNull(newWrapMetadata);
-#else
-            if (newWrapMetadata == null)
-            {
-                throw new ArgumentNullException(nameof(newWrapMetadata));
-            }
-#endif
+            ArgumentValidation.ThrowIfNull(newWrapMetadata);
 
             return TaskHelper.RunInlineIfNeededAsync(() =>
                 this.dataEncryptionKeyContainerCore.RewrapDataEncryptionKeyAsync(id, newWrapMetadata, encryptionAlgorithm, requestOptions, cancellationToken));

@@ -45,14 +45,7 @@ internal static DecryptableFeedResponse<T> CreateResponse(
             ResponseMessage responseMessage,
             IReadOnlyCollection<T> resource)
         {
-#if NET8_0_OR_GREATER
-            ArgumentNullException.ThrowIfNull(responseMessage);
-#else
-            if (responseMessage == null)
-            {
-                throw new ArgumentNullException(nameof(responseMessage));
-            }
-#endif
+            ArgumentValidation.ThrowIfNull(responseMessage);
 
             using (responseMessage)
             {

diff --git a/Microsoft.Azure.Cosmos.Encryption.Custom/src/Diagnostics/EncryptionDiagnostics.cs b/Microsoft.Azure.Cosmos.Encryption.Custom/src/Diagnostics/EncryptionDiagnostics.cs
@@ -0,0 +1,11 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+namespace Microsoft.Azure.Cosmos.Encryption.Custom
+{
+    internal static class EncryptionDiagnostics
+    {
+        internal const string ScopeEncryptModeSelectionPrefix = "EncryptionProcessor.Encrypt.Mde.";
+        internal const string ScopeDecryptModeSelectionPrefix = "EncryptionProcessor.Decrypt.Mde.";
+    }
+}