Add message context to serialization interfaces

Directory.Build.props

@@ -9,6 +9,9 @@
 
         <!-- Do not package by default to avoid accidental publishing -->
         <IsPackable>false</IsPackable>
+
+        <!-- NuGet version used by ABI compatibility checks -->
+        <NatsAbiCheckVersion>2.6.0</NatsAbiCheckVersion>
     </PropertyGroup>
 
     <ItemGroup>

NATS.Net.sln

@@ -158,6 +158,8 @@ Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "NATS.Client.CheckAbiTransie
 EndProject
 Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "NATS.Client.CheckAbi", "tests\NATS.Client.CheckAbi\NATS.Client.CheckAbi.csproj", "{DA8E9EC2-990C-4DA2-A3F3-7E7796C71DD8}"
 EndProject
+Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "NATS.Client.CheckAbiOld", "tests\NATS.Client.CheckAbiOld\NATS.Client.CheckAbiOld.csproj", "{93E92151-28BF-4176-92F4-65E1227C3269}"
+EndProject
 Global
 	GlobalSection(SolutionConfigurationPlatforms) = preSolution
 		Debug|Any CPU = Debug|Any CPU
@@ -372,6 +374,10 @@ Global
 		{DA8E9EC2-990C-4DA2-A3F3-7E7796C71DD8}.Debug|Any CPU.Build.0 = Debug|Any CPU
 		{DA8E9EC2-990C-4DA2-A3F3-7E7796C71DD8}.Release|Any CPU.ActiveCfg = Release|Any CPU
 		{DA8E9EC2-990C-4DA2-A3F3-7E7796C71DD8}.Release|Any CPU.Build.0 = Release|Any CPU
+		{93E92151-28BF-4176-92F4-65E1227C3269}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
+		{93E92151-28BF-4176-92F4-65E1227C3269}.Debug|Any CPU.Build.0 = Debug|Any CPU
+		{93E92151-28BF-4176-92F4-65E1227C3269}.Release|Any CPU.ActiveCfg = Release|Any CPU
+		{93E92151-28BF-4176-92F4-65E1227C3269}.Release|Any CPU.Build.0 = Release|Any CPU
 	EndGlobalSection
 	GlobalSection(SolutionProperties) = preSolution
 		HideSolutionNode = FALSE
@@ -432,6 +438,7 @@ Global
 		{71C2BC96-C104-4748-B6CF-E74BAC4E709D} = {4827B3EC-73D8-436D-AE2A-5E29AC95FD0C}
 		{21649885-B40A-4F6B-9286-A7E378276C3E} = {C526E8AB-739A-48D7-8FC4-048978C9B650}
 		{DA8E9EC2-990C-4DA2-A3F3-7E7796C71DD8} = {C526E8AB-739A-48D7-8FC4-048978C9B650}
+		{93E92151-28BF-4176-92F4-65E1227C3269} = {C526E8AB-739A-48D7-8FC4-048978C9B650}
 	EndGlobalSection
 	GlobalSection(ExtensibilityGlobals) = postSolution
 		SolutionGuid = {8CBB7278-D093-448E-B3DE-B5991209A1AA}

src/NATS.Client.Abstractions/INatsHeaders.cs

@@ -0,0 +1,10 @@
+using Microsoft.Extensions.Primitives;
+
+namespace NATS.Client.Core;
+
+/// <summary>
+/// Represents NATS message headers as a dictionary of string keys and <see cref="StringValues"/> values.
+/// </summary>
+public interface INatsHeaders : IDictionary<string, StringValues>
+{
+}

src/NATS.Client.Abstractions/INatsSerialize.cs

@@ -44,9 +44,89 @@ public interface INatsDeserialize<out T>
     T? Deserialize(in ReadOnlySequence<byte> buffer);
 }
 
+/// <summary>
+/// Extended serializer interface with access to message context during serialization.
+/// </summary>
+/// <typeparam name="T">Serialized object type</typeparam>
+public interface INatsSerializeWithContext<in T> : INatsSerialize<T>
+{
+    /// <summary>
+    /// Serialize value to buffer with message context.
+    /// </summary>
+    /// <param name="bufferWriter">Buffer to write the serialized data.</param>
+    /// <param name="value">Object to be serialized.</param>
+    /// <param name="context">Message envelope metadata.</param>
+    void Serialize(IBufferWriter<byte> bufferWriter, T value, in NatsMsgContext context);
+}
+
+/// <summary>
+/// Extended deserializer interface with access to message context during deserialization.
+/// </summary>
+/// <typeparam name="T">Deserialized object type</typeparam>
+public interface INatsDeserializeWithContext<out T> : INatsDeserialize<T>
+{
+    /// <summary>
+    /// Deserialize value from buffer with message context.
+    /// </summary>
+    /// <param name="buffer">Buffer with the serialized data.</param>
+    /// <param name="context">Message envelope metadata.</param>
+    /// <returns>Deserialized object</returns>
+    T? Deserialize(in ReadOnlySequence<byte> buffer, in NatsMsgContext context);
+}
+
+/// <summary>
+/// Combined context-aware serializer interface that supports both serialization and deserialization
+/// with access to message context.
+/// </summary>
+/// <typeparam name="T">Object type</typeparam>
+public interface INatsSerializerWithContext<T> : INatsSerializeWithContext<T>, INatsDeserializeWithContext<T>
+{
+}
+
+/// <summary>
+/// Registry for serializers and deserializers.
+/// </summary>
 public interface INatsSerializerRegistry
 {
     INatsSerialize<T> GetSerializer<T>();
 
     INatsDeserialize<T> GetDeserializer<T>();
 }
+
+/// <summary>
+/// Extension methods to support context-aware serialization with fallback to standard serialization.
+/// </summary>
+/// <remarks>
+/// Context-aware serializers that want to mutate <see cref="NatsMsgContext.Headers"/> must check
+/// for null before doing so: the library passes whatever headers the caller supplied to
+/// <c>PublishAsync</c>, and does not allocate a new <see cref="INatsHeaders"/> for serializers that
+/// opt in to context. Callers who want header-mutating behavior must pass a non-null
+/// headers instance to <c>PublishAsync</c>.
+/// </remarks>
+public static class NatsSerializationExtensions
+{
+    /// <summary>
+    /// Serializes the value with message context, falling back to standard serialization if not supported.
+    /// </summary>
+    public static void Serialize<T>(this INatsSerialize<T> serializer, IBufferWriter<byte> bufferWriter, T value, in NatsMsgContext context)
+    {
+        if (serializer is INatsSerializeWithContext<T> withContext)
+        {
+            withContext.Serialize(bufferWriter, value, in context);
+            return;
+        }
+
+        serializer.Serialize(bufferWriter, value);
+    }
+
+    /// <summary>
+    /// Deserializes the value with message context, falling back to standard deserialization if not supported.
+    /// </summary>
+    public static T? Deserialize<T>(this INatsDeserialize<T> deserializer, in ReadOnlySequence<byte> buffer, in NatsMsgContext context)
+    {
+        if (deserializer is INatsDeserializeWithContext<T> withContext)
+            return withContext.Deserialize(buffer, in context);
+
+        return deserializer.Deserialize(buffer);
+    }
+}

src/NATS.Client.Abstractions/NATS.Client.Abstractions.csproj

@@ -8,6 +8,7 @@
 
   <!--  Dependencies for all -->
   <ItemGroup>
+    <PackageReference Include="Microsoft.Extensions.Primitives" Version="6.0.0"/>
     <PackageReference Include="Microsoft.VisualStudio.Threading.Analyzers" Version="17.10.48">
       <PrivateAssets>all</PrivateAssets>
       <IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>

src/NATS.Client.Abstractions/NatsMsgContext.cs

@@ -0,0 +1,32 @@
+namespace NATS.Client.Core;
+
+/// <summary>
+/// Provides message envelope metadata available during serialization and deserialization.
+/// </summary>
+/// <remarks>
+/// Use the constructor to create instances. A <c>default(NatsMsgContext)</c> value has a null
+/// <see cref="Subject"/> and is not a valid context; the library never passes such a value to
+/// user code.
+/// </remarks>
+public readonly struct NatsMsgContext
+{
+    /// <summary>Creates a new <see cref="NatsMsgContext"/>.</summary>
+    /// <param name="subject">Subject the message was published to.</param>
+    /// <param name="replyTo">Optional reply-to subject.</param>
+    /// <param name="headers">Optional message headers.</param>
+    public NatsMsgContext(string subject, string? replyTo = null, INatsHeaders? headers = null)
+    {
+        Subject = subject ?? throw new ArgumentNullException(nameof(subject));
+        ReplyTo = replyTo;
+        Headers = headers;
+    }
+
+    /// <summary>Subject the message was published to.</summary>
+    public string Subject { get; }
+
+    /// <summary>Optional reply-to subject.</summary>
+    public string? ReplyTo { get; }
+
+    /// <summary>Optional message headers. May be null if the caller did not supply headers.</summary>
+    public INatsHeaders? Headers { get; }
+}

src/NATS.Client.Core/Commands/CommandWriter.cs

@@ -334,12 +334,12 @@ public ValueTask PublishAsync<T>(string subject, T? value, NatsHeaders? headers,
 
         try
         {
+            if (value != null)
+                serializer.Serialize(payloadBuffer, value, new NatsMsgContext(subject, replyTo, headers));
+
             if (headers != null)
                 _headerWriter.Write(headersBuffer!, headers);
 
-            if (value != null)
-                serializer.Serialize(payloadBuffer, value);
-
             var size = payloadBuffer.WrittenMemory.Length + (headersBuffer?.WrittenMemory.Length ?? 0);
             if (_connection.ServerInfo is { } info && size > info.MaxPayload)
             {

src/NATS.Client.Core/Internal/Telemetry.cs

@@ -104,11 +104,7 @@ public static void AddTraceContextHeaders(Activity? activity, ref NatsHeaders? h
                     return;
                 }
 
-                // There are cases where headers reused publicly (e.g. JetStream publish retry)
-                // there may also be cases where application can reuse the same header
-                // in which case we should still be able to overwrite headers with telemetry fields
-                // even though headers would be set to readonly before being passed down in publish methods.
-                headers.SetOverrideReadOnly(fieldName, fieldValue);
+                headers[fieldName] = fieldValue;
             });
     }
 

src/NATS.Client.Core/NatsConnection.Publish.cs

@@ -19,7 +19,6 @@ public ValueTask PublishAsync(string subject, NatsHeaders? headers = default, st
             Telemetry.AddTraceContextHeaders(activity, ref headers);
             try
             {
-                headers?.SetReadOnly();
                 return ConnectionState != NatsConnectionState.Open
                     ? ConnectAndPublishAsync(subject, default, headers, replyTo, NatsRawSerializer<byte[]>.Default, cancellationToken)
                     : CommandWriter.PublishAsync(subject, default, headers, replyTo, NatsRawSerializer<byte[]>.Default, cancellationToken);
@@ -31,7 +30,6 @@ public ValueTask PublishAsync(string subject, NatsHeaders? headers = default, st
             }
         }
 
-        headers?.SetReadOnly();
         return ConnectionState != NatsConnectionState.Open
             ? ConnectAndPublishAsync(subject, default, headers, replyTo, NatsRawSerializer<byte[]>.Default, cancellationToken)
             : CommandWriter.PublishAsync(subject, default, headers, replyTo, NatsRawSerializer<byte[]>.Default, cancellationToken);
@@ -53,7 +51,6 @@ public ValueTask PublishAsync<T>(string subject, T? data, NatsHeaders? headers =
             try
             {
                 serializer ??= Opts.SerializerRegistry.GetSerializer<T>();
-                headers?.SetReadOnly();
                 return ConnectionState != NatsConnectionState.Open
                     ? ConnectAndPublishAsync(subject, data, headers, replyTo, serializer, cancellationToken)
                     : CommandWriter.PublishAsync(subject, data, headers, replyTo, serializer, cancellationToken);
@@ -66,7 +63,6 @@ public ValueTask PublishAsync<T>(string subject, T? data, NatsHeaders? headers =
         }
 
         serializer ??= Opts.SerializerRegistry.GetSerializer<T>();
-        headers?.SetReadOnly();
         return ConnectionState != NatsConnectionState.Open
             ? ConnectAndPublishAsync(subject, data, headers, replyTo, serializer, cancellationToken)
             : CommandWriter.PublishAsync(subject, data, headers, replyTo, serializer, cancellationToken);

src/NATS.Client.Core/NatsHeaders.cs

@@ -10,12 +10,18 @@ namespace NATS.Client.Core;
 /// <summary>
 /// Represents a wrapper for RequestHeaders and ResponseHeaders.
 /// </summary>
+/// <remarks>
+/// Not thread-safe. Do not share a single <see cref="NatsHeaders"/> instance across concurrent
+/// publishes: the writer may read and context-aware serializers may mutate the dictionary while
+/// it is being written to the wire. Construct a fresh instance per publish, or wait for one
+/// publish to complete before reusing the same instance.
+/// </remarks>
 [SuppressMessage("StyleCop.CSharp.OrderingRules", "SA1201:Elements should appear in the correct order", Justification = "Keep class format as is for reference")]
 [SuppressMessage("StyleCop.CSharp.LayoutRules", "SA1504:All accessors should be single-line or multi-line", Justification = "Keep class format as is for reference")]
 [SuppressMessage("StyleCop.CSharp.LayoutRules", "SA1516:Elements should be separated by blank line", Justification = "Keep class format as is for reference")]
 [SuppressMessage("StyleCop.CSharp.LayoutRules", "SA1513:Closing brace should be followed by blank line", Justification = "Keep class format as is for reference")]
 [SuppressMessage("StyleCop.CSharp.OrderingRules", "SA1214:Readonly fields should appear before non-readonly fields", Justification = "Keep class format as is for reference")]
-public class NatsHeaders : IDictionary<string, StringValues>
+public class NatsHeaders : INatsHeaders
 {
     public enum Messages
     {
@@ -70,8 +76,6 @@ public enum Messages
     private static readonly IEnumerator<KeyValuePair<string, StringValues>> EmptyIEnumeratorType = default(Enumerator);
     private static readonly IEnumerator EmptyIEnumerator = default(Enumerator);
 
-    private int _readonly = 0;
-
     public int Version => 1;
 
     public int Code { get; internal set; }
@@ -144,7 +148,6 @@ public StringValues this[string key]
             {
                 throw new ArgumentNullException(nameof(key));
             }
-            ThrowIfReadOnly();
 
             if (value.Count == 0)
             {
@@ -161,11 +164,7 @@ public StringValues this[string key]
     StringValues IDictionary<string, StringValues>.this[string key]
     {
         get { return this[key]; }
-        set
-        {
-            ThrowIfReadOnly();
-            this[key] = value;
-        }
+        set { this[key] = value; }
     }
 
     /// <summary>
@@ -175,10 +174,9 @@ StringValues IDictionary<string, StringValues>.this[string key]
     public int Count => Store?.Count ?? 0;
 
     /// <summary>
-    /// Gets a value that indicates whether the <see cref="NatsHeaders" /> is in read-only mode.
+    /// Gets a value that indicates whether the <see cref="NatsHeaders" /> is read-only. Always false.
     /// </summary>
-    /// <returns>true if the <see cref="NatsHeaders" /> is in read-only mode; otherwise, false.</returns>
-    public bool IsReadOnly => Volatile.Read(ref _readonly) == 1;
+    public bool IsReadOnly => false;
 
     /// <summary>
     /// Gets the collection of HTTP header names in this instance.
@@ -222,7 +220,6 @@ public void Add(KeyValuePair<string, StringValues> item)
         {
             throw new ArgumentException("The key is null");
         }
-        ThrowIfReadOnly();
         EnsureStore(1);
         Store.Add(item.Key, item.Value);
     }
@@ -238,7 +235,6 @@ public void Add(string key, StringValues value)
         {
             throw new ArgumentNullException(nameof(key));
         }
-        ThrowIfReadOnly();
         EnsureStore(1);
         Store.Add(key, value);
     }
@@ -248,7 +244,6 @@ public void Add(string key, StringValues value)
     /// </summary>
     public void Clear()
     {
-        ThrowIfReadOnly();
         Store?.Clear();
     }
 
@@ -322,7 +317,6 @@ public int GetBytesLength(Encoding? encoding = null)
     /// <returns>true if the specified object was removed from the collection; otherwise, false.</returns>
     public bool Remove(KeyValuePair<string, StringValues> item)
     {
-        ThrowIfReadOnly();
         if (Store == null)
         {
             return false;
@@ -342,7 +336,6 @@ public bool Remove(KeyValuePair<string, StringValues> item)
     /// <returns>true if the specified object was removed from the collection; otherwise, false.</returns>
     public bool Remove(string key)
     {
-        ThrowIfReadOnly();
         if (Store == null)
         {
             return false;
@@ -426,34 +419,6 @@ IEnumerator IEnumerable.GetEnumerator()
         return Store.GetEnumerator();
     }
 
-    internal void SetReadOnly() => Interlocked.Exchange(ref _readonly, 1);
-
-    internal void SetOverrideReadOnly(string key, StringValues value)
-    {
-        if (key == null)
-        {
-            throw new ArgumentNullException(nameof(key));
-        }
-
-        if (value.Count == 0)
-        {
-            Store?.Remove(key);
-        }
-        else
-        {
-            EnsureStore(1);
-            Store[key] = value;
-        }
-    }
-
-    private void ThrowIfReadOnly()
-    {
-        if (IsReadOnly)
-        {
-            throw new InvalidOperationException("The response headers cannot be modified because the response has already started.");
-        }
-    }
-
     /// <summary>
     /// Enumerates a <see cref="NatsHeaders"/>.
     /// </summary>