[C#] Remote pubsub

cs/remote/benchmark/FASTER.benchmark/Types.cs

@@ -69,5 +69,12 @@ public void UpsertCompletionCallback(ref Key key, ref Value value, Empty ctx)
         public void SubscribeKVCallback(ref Key key, ref Input input, ref Output output, Empty ctx, Status status)
         {
         }
+
+        public void PublishCompletionCallback(ref Key key, ref Value value, Empty ctx)
+        {
+        }
+        public void SubscribeCallback(ref Key key, ref Value value, Empty ctx)
+        {
+        }
     }
 }

cs/remote/samples/FixedLenServer/Program.cs

@@ -37,11 +37,14 @@ static void FixedLenServer(string[] args)
             var store = new FasterKV<Key, Value>(indexSize, logSettings, checkpointSettings);
             if (opts.Recover) store.Recover();
 
-            var broker = new SubscribeKVBroker<Key, Value, IKeySerializer<Key>>(new FixedLenKeySerializer<Key>());
+            // Create a broker for pub-sub of key-value pairs in remote FASTER instance
+            var kvBroker = new SubscribeKVBroker<Key, Value, IKeySerializer<Key>>(new FixedLenKeySerializer<Key>());
+            // Create a broker for pub-sub of key-value pairs
+            var broker = new SubscribeBroker<Key, Value, IKeySerializer<Key>>(new FixedLenKeySerializer<Key>());
 
             // This fixed-length session provider can be used with compatible clients such as FixedLenClient and FASTER.benchmark
             // Uses FixedLenSerializer as our in-built serializer for blittable (fixed length) types
-            var provider = new FasterKVProvider<Key, Value, Input, Output, Functions, FixedLenSerializer<Key, Value, Input, Output>>(store, e => new Functions(), broker);
+            var provider = new FasterKVProvider<Key, Value, Input, Output, Functions, FixedLenSerializer<Key, Value, Input, Output>>(store, e => new Functions(), kvBroker, broker);
 
             // Create server
             var server = new FasterServer(opts.Address, opts.Port);

cs/remote/samples/VarLenClient/CustomTypeFunctions.cs

@@ -33,6 +33,8 @@ public override void SubscribeKVCallback(ref CustomType key, ref CustomType inpu
         {
         }
 
-
+        public override void SubscribeCallback(ref CustomType key, ref CustomType value, byte ctx)
+        {
+        }
     }
 }

cs/remote/samples/VarLenClient/CustomTypeSamples.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT license.
 
 using System;
+using System.Threading;
 using System.Threading.Tasks;
 using FASTER.client;
 using FASTER.common;
@@ -28,6 +29,9 @@ public void Run(string ip, int port)
             // Samples using sync client API
             SyncVarLenSamples(session);
 
+            // Samples using sync client API
+            SyncVarLenSubscriptionKVSamples(session, subSession);
+
             // Samples using sync client API
             SyncVarLenSubscriptionSamples(session, subSession);
 
@@ -54,7 +58,7 @@ void SyncVarLenSamples(ClientSession<CustomType, CustomType, CustomType, CustomT
             session.CompletePending(true);
         }
 
-        void SyncVarLenSubscriptionSamples(ClientSession<CustomType, CustomType, CustomType, CustomType, byte, CustomTypeFunctions, FixedLenSerializer<CustomType, CustomType, CustomType, CustomType>> session,
+        void SyncVarLenSubscriptionKVSamples(ClientSession<CustomType, CustomType, CustomType, CustomType, byte, CustomTypeFunctions, FixedLenSerializer<CustomType, CustomType, CustomType, CustomType>> session,
                                             ClientSession<CustomType, CustomType, CustomType, CustomType, byte, CustomTypeFunctions, FixedLenSerializer<CustomType, CustomType, CustomType, CustomType>> session2)
         {
             session2.SubscribeKV(new CustomType(23));
@@ -78,6 +82,29 @@ void SyncVarLenSubscriptionSamples(ClientSession<CustomType, CustomType, CustomT
             System.Threading.Thread.Sleep(1000);
         }
 
+        void SyncVarLenSubscriptionSamples(ClientSession<CustomType, CustomType, CustomType, CustomType, byte, CustomTypeFunctions, FixedLenSerializer<CustomType, CustomType, CustomType, CustomType>> session,
+                                            ClientSession<CustomType, CustomType, CustomType, CustomType, byte, CustomTypeFunctions, FixedLenSerializer<CustomType, CustomType, CustomType, CustomType>> session2)
+        {
+            session2.Subscribe(new CustomType(23));
+            session2.CompletePending(true);
+
+            session2.Subscribe(new CustomType(24));
+            session2.CompletePending(true);
+
+            session2.PSubscribe(new CustomType(25));
+            session2.CompletePending(true);
+
+            session.Publish(new CustomType(23), new CustomType(2300));
+            session.CompletePending(true);
+
+            session.Publish(new CustomType(24), new CustomType(2400));
+            session.CompletePending(true);
+
+            session.Publish(new CustomType(25), new CustomType(2500));
+            session.CompletePending(true);
+
+            System.Threading.Thread.Sleep(1000);
+        }
 
         async Task AsyncVarLenSamples(ClientSession<CustomType, CustomType, CustomType, CustomType, byte, CustomTypeFunctions, FixedLenSerializer<CustomType, CustomType, CustomType, CustomType>> session)
         {

cs/remote/samples/VarLenServer/Program.cs

@@ -37,10 +37,13 @@ static void VarLenServer(string[] args)
             var store = new FasterKV<SpanByte, SpanByte>(indexSize, logSettings, checkpointSettings);
             if (opts.Recover) store.Recover();
 
-            var broker = new SubscribeKVBroker<SpanByte, SpanByte, IKeySerializer<SpanByte>>(new SpanByteKeySerializer());
+            // Create a broker for pub-sub of key value-pairs in remote FASTER instance
+            var kvBroker = new SubscribeKVBroker<SpanByte, SpanByte, IKeySerializer<SpanByte>>(new SpanByteKeySerializer());
+            // Create a broker for pub-sub of key-value pairs
+            var broker = new SubscribeBroker<SpanByte, SpanByte, IKeySerializer<SpanByte>>(new SpanByteKeySerializer());
 
             // This variable-length session provider can be used with compatible clients such as VarLenClient
-            var provider = new SpanByteFasterKVProvider(store, broker);
+            var provider = new SpanByteFasterKVProvider(store, kvBroker, broker);
 
             // Create server
             var server = new FasterServer(opts.Address, opts.Port);

cs/remote/src/FASTER.client/CallbackFunctionsBase.cs

@@ -25,5 +25,9 @@ public virtual void RMWCompletionCallback(ref Key key, ref Input input, ref Outp
         public virtual void UpsertCompletionCallback(ref Key key, ref Value value, Context ctx) { }
         /// <inheritdoc/>
         public virtual void SubscribeKVCallback(ref Key key, ref Input input, ref Output output, Context ctx, Status status) { }
+        /// <inheritdoc/>
+        public virtual void PublishCompletionCallback(ref Key key, ref Value value, Context ctx) { }
+        /// <inheritdoc/>
+        public virtual void SubscribeCallback(ref Key key, ref Value value, Context ctx) { }
     }
 }

cs/remote/src/FASTER.client/ClientSession.cs

@@ -45,6 +45,7 @@ public unsafe sealed partial class ClientSession<Key, Value, Input, Output, Cont
 
         readonly ElasticCircularBuffer<(Key, Value, Context)> upsertQueue;
         readonly ElasticCircularBuffer<(Key, Input, Output, Context)> readrmwQueue;
+        readonly ElasticCircularBuffer<(Key, Value, Context)> pubsubQueue;
         readonly ElasticCircularBuffer<TaskCompletionSource<(Status, Output)>> tcsQueue;
 
         /// <summary>
@@ -69,6 +70,7 @@ public ClientSession(string address, int port, Functions functions, WireFormat w
 
             upsertQueue = new ElasticCircularBuffer<(Key, Value, Context)>();
             readrmwQueue = new ElasticCircularBuffer<(Key, Input, Output, Context)>();
+            pubsubQueue = new ElasticCircularBuffer<(Key, Value, Context)>();
             tcsQueue = new ElasticCircularBuffer<TaskCompletionSource<(Status, Output)>>();
 
             numPendingBatches = 0;
@@ -238,7 +240,38 @@ public void SubscribeKV(Key key, Input input = default, Context userContext = de
         public void PSubscribeKV(Key prefix, Input input = default, Context userContext = default, long serialNo = 0)
             => InternalSubscribeKV(MessageType.PSubscribeKV, ref prefix, ref input, userContext, serialNo);
 
+        /// <summary>
+        /// Upsert operation
+        /// </summary>
+        /// <param name="key">Key</param>
+        /// <param name="desiredValue">Desired value</param>
+        /// <param name="userContext">User context</param>
+        /// <param name="serialNo">Serial number</param>
+        /// <returns>Status of operation</returns>
+        public Status Publish(Key key, Value desiredValue, Context userContext = default, long serialNo = 0)
+            => InternalPublish(MessageType.Publish, ref key, ref desiredValue, userContext, serialNo);
 
+        /// <summary>
+        /// SubscribeKV operation
+        /// </summary>
+        /// <param name="key">Key</param>
+        /// <param name="input">Input</param>
+        /// <param name="userContext">User context</param>
+        /// <param name="serialNo">Serial number</param>
+        /// <returns>Status of operation</returns>
+        public void Subscribe(Key key, Context userContext = default, long serialNo = 0)
+            => InternalSubscribe(MessageType.Subscribe, ref key, userContext, serialNo);
+
+        /// <summary>
+        /// PSubscribe operation
+        /// </summary>
+        /// <param name="prefix">Key</param>
+        /// <param name="input">Input</param>
+        /// <param name="userContext">User context</param>
+        /// <param name="serialNo">Serial number</param>
+        /// <returns>Status of operation</returns>
+        public void PSubscribe(Key prefix, Context userContext = default, long serialNo = 0)
+            => InternalSubscribe(MessageType.PSubscribe, ref prefix, userContext, serialNo);
 
         /// <summary>
         /// Flush current buffer of outgoing messages. Does not wait for responses.
@@ -307,6 +340,7 @@ public void Dispose(bool completePending)
 
 
         int lastSeqNo = -1;
+        readonly Dictionary<int, (Key, Value, Context)> pubsubPendingContext = new();
         readonly Dictionary<int, (Key, Input, Output, Context)> readRmwPendingContext = new();
         readonly Dictionary<int, TaskCompletionSource<(Status, Output)>> readRmwPendingTcs = new();
 
@@ -481,6 +515,56 @@ internal void ProcessReplies(byte[] buf, int offset)
                                 }
                                 break;
                             }
+                        case MessageType.Publish:
+                            {
+                                var status = ReadStatus(ref src);
+                                (Key, Value, Context) result = upsertQueue.Dequeue();
+                                functions.PublishCompletionCallback(ref result.Item1, ref result.Item2, result.Item3);
+                                break;
+                            }
+                        case MessageType.Subscribe:
+                            {
+                                var status = ReadStatus(ref src);
+                                var p = hrw.ReadPendingSeqNo(ref src);
+                                if (status == Status.OK)
+                                {
+                                    pubsubPendingContext.TryGetValue(p, out var result);
+                                    result.Item2 = serializer.ReadValue(ref src);
+                                    functions.SubscribeCallback(ref result.Item1, ref result.Item2, result.Item3);
+                                }
+                                else if (status == Status.PENDING)
+                                {
+                                    var result = pubsubQueue.Dequeue();
+                                    pubsubPendingContext.Add(p, result);
+                                }
+                                else
+                                {
+                                    throw new Exception("Unexpected status of SubscribeKV");
+                                }
+                                break;
+                            }
+                        case MessageType.PSubscribe:
+                            {
+                                var status = ReadStatus(ref src);
+                                var p = hrw.ReadPendingSeqNo(ref src);
+                                if (status == Status.OK)
+                                {
+                                    pubsubPendingContext.TryGetValue(p, out var result);
+                                    result.Item1 = serializer.ReadKey(ref src);
+                                    result.Item2 = serializer.ReadValue(ref src);
+                                    functions.SubscribeCallback(ref result.Item1, ref result.Item2, result.Item3);
+                                }
+                                else if (status == Status.PENDING)
+                                {
+                                    var result = pubsubQueue.Dequeue();
+                                    pubsubPendingContext.Add(p, result);
+                                }
+                                else
+                                {
+                                    throw new Exception("Unexpected status of SubscribeKV");
+                                }
+                                break;
+                            }
                         case MessageType.PendingResult:
                             {
                                 HandlePending(ref src);
@@ -667,12 +751,58 @@ private unsafe Status InternalSubscribeKV(MessageType messageType, ref Key key,
                 byte* end = sendObject.obj.bufferPtr + bufferSize;
                 byte* curr = offset;
                 if (hrw.Write(messageType, ref curr, (int)(end - curr)))
-                    if (serializer.Write(ref key, ref curr, (int)(end - curr)))
-                        if (serializer.Write(ref input, ref curr, (int)(end - curr)))
+                    if (hrw.Write(serialNo, ref curr, (int)(end - curr)))
+                        if (serializer.Write(ref key, ref curr, (int)(end - curr)))
+                            if (serializer.Write(ref input, ref curr, (int)(end - curr)))
+                            {
+                                numMessages++;
+                                offset = curr;
+                                readrmwQueue.Enqueue((key, input, default, userContext));
+                                return Status.PENDING;
+                            }
+                Flush();
+            }
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private unsafe Status InternalPublish(MessageType messageType, ref Key key, ref Value desiredValue, Context userContext = default, long serialNo = 0)
+        {
+            Debug.Assert(!subscriptionSession);
+
+            while (true)
+            {
+                byte* end = sendObject.obj.bufferPtr + bufferSize;
+                byte* curr = offset;
+                if (hrw.Write(messageType, ref curr, (int)(end - curr)))
+                    if (hrw.Write(serialNo, ref curr, (int)(end - curr)))
+                        if (serializer.Write(ref key, ref curr, (int)(end - curr)))
+                            if (serializer.Write(ref desiredValue, ref curr, (int)(end - curr)))
+                            {
+                                numMessages++;
+                                offset = curr;
+                                upsertQueue.Enqueue((key, desiredValue, userContext));
+                                return Status.PENDING;
+                            }
+                Flush();
+            }
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private unsafe Status InternalSubscribe(MessageType messageType, ref Key key, Context userContext = default, long serialNo = 0)
+        {
+            subscriptionSession = true;
+
+            while (true)
+            {
+                byte* end = sendObject.obj.bufferPtr + bufferSize;
+                byte* curr = offset;
+                if (hrw.Write(messageType, ref curr, (int)(end - curr)))
+                    if (hrw.Write(serialNo, ref curr, (int)(end - curr)))
+                        if (serializer.Write(ref key, ref curr, (int)(end - curr)))
                         {
                             numMessages++;
                             offset = curr;
-                            readrmwQueue.Enqueue((key, input, default, userContext));
+                            pubsubQueue.Enqueue((key, default, userContext));
                             return Status.PENDING;
                         }
                 Flush();

cs/remote/src/FASTER.client/FixedLenSerializer.cs

@@ -62,6 +62,14 @@ public Key ReadKey(ref byte* src)
             return Unsafe.AsRef<Key>(_src);
         }
 
+        /// <inheritdoc />
+        public Value ReadValue(ref byte* src)
+        {
+            var _src = src;
+            src += Unsafe.SizeOf<Value>();
+            return Unsafe.AsRef<Value>(_src);
+        }
+
         /// <inheritdoc />
         public Output ReadOutput(ref byte* src)
         {

cs/remote/src/FASTER.client/ICallbackFunctions.cs

@@ -59,5 +59,21 @@ public interface ICallbackFunctions<Key, Value, Input, Output, Context>
         /// <param name="ctx"></param>
         /// <param name="status"></param>
         void SubscribeKVCallback(ref Key key, ref Input input, ref Output output, Context ctx, Status status);
+
+        /// <summary>
+        /// Publish completion
+        /// </summary>
+        /// <param name="key"></param>
+        /// <param name="value"></param>
+        /// <param name="ctx"></param>
+        void PublishCompletionCallback(ref Key key, ref Value value, Context ctx);
+
+        /// <summary>
+        /// Subscribe callback
+        /// </summary>
+        /// <param name="key"></param>
+        /// <param name="value"></param>
+        /// <param name="ctx"></param>
+        void SubscribeCallback(ref Key key, ref Value value, Context ctx);
     }
 }

cs/remote/src/FASTER.client/MemoryFunctionsBase.cs

@@ -26,5 +26,9 @@ public virtual void RMWCompletionCallback(ref ReadOnlyMemory<T> key, ref ReadOnl
         public virtual void UpsertCompletionCallback(ref ReadOnlyMemory<T> key, ref ReadOnlyMemory<T> value, byte ctx) { }
         /// <inheritdoc />
         public virtual void SubscribeKVCallback(ref ReadOnlyMemory<T> key, ref ReadOnlyMemory<T> input, ref (IMemoryOwner<T>, int) output, byte ctx, Status status) { }
+        /// <inheritdoc/>
+        public virtual void PublishCompletionCallback(ref ReadOnlyMemory<T> key, ref ReadOnlyMemory<T> value, byte ctx) { }
+        /// <inheritdoc/>
+        public virtual void SubscribeCallback(ref ReadOnlyMemory<T> key, ref ReadOnlyMemory<T> value, byte ctx) { }
     }
 }

cs/remote/src/FASTER.client/MemoryParameterSerializer.cs

@@ -37,6 +37,17 @@ public ReadOnlyMemory<T> ReadKey(ref byte* src)
             return (mem.Memory.Slice(0, len));
         }
 
+        /// <inheritdoc />
+        public ReadOnlyMemory<T> ReadValue(ref byte* src)
+        {
+            var len = (*(int*)src) / sizeof(T);
+            var mem = memoryPool.Rent(len);
+            new ReadOnlySpan<byte>(src + sizeof(int), (*(int*)src)).CopyTo(
+               MemoryMarshal.Cast<T, byte>(mem.Memory.Span));
+            src += sizeof(int) + (*(int*)src);
+            return (mem.Memory.Slice(0, len));
+        }
+
         /// <inheritdoc />
         public (IMemoryOwner<T>, int) ReadOutput(ref byte* src)
         {