C#: Custom command for standalone client. (#3335)

* Custom command for standalone.

Signed-off-by: Yury-Fridlyand <[email protected]>

Author: Yury-Fridlyand

csharp/lib/BaseClient.cs

@@ -29,41 +29,40 @@ protected BaseClient(BaseClientConfiguration config)
         }
     }
 
-    protected async Task<string?> Command(IntPtr[] args, int argsCount, RequestType requestType)
+    protected async Task<T> Command<T>(string[] arguments, RequestType requestType) where T : class?
     {
+        IntPtr[] args = _arrayPool.Rent(arguments.Length);
+        for (int i = 0; i < arguments.Length; i++)
+        {
+            args[i] = Marshal.StringToHGlobalAnsi(arguments[i]);
+        }
         // We need to pin the array in place, in order to ensure that the GC doesn't move it while the operation is running.
         GCHandle pinnedArray = GCHandle.Alloc(args, GCHandleType.Pinned);
         IntPtr pointer = pinnedArray.AddrOfPinnedObject();
-        Message<string> message = _messageContainer.GetMessageForCall(args, argsCount);
-        CommandFfi(_clientPointer, (ulong)message.Index, (int)requestType, pointer, (uint)argsCount);
-        string? result = await message;
+        Message message = _messageContainer.GetMessageForCall<T>(args);
+        CommandFfi(_clientPointer, (ulong)message.Index, (int)requestType, pointer, (uint)arguments.Length);
+        for (int i = 0; i < arguments.Length; i++)
+        {
+            Marshal.FreeHGlobal(args[i]);
+        }
         pinnedArray.Free();
-        return result;
-    }
-
-    public async Task<string?> Set(string key, string value)
-    {
-        IntPtr[] args = _arrayPool.Rent(2);
-        args[0] = Marshal.StringToHGlobalAnsi(key);
-        args[1] = Marshal.StringToHGlobalAnsi(value);
-        string? result = await Command(args, 2, RequestType.Set);
         _arrayPool.Return(args);
-        return result;
+#pragma warning disable CS8603 // Possible null reference return.
+        return await message as T;
+#pragma warning restore CS8603 // Possible null reference return.
     }
 
+    public async Task<string> Set(string key, string value)
+        => await Command<string>([key, value], RequestType.Set);
+
     public async Task<string?> Get(string key)
-    {
-        IntPtr[] args = _arrayPool.Rent(1);
-        args[0] = Marshal.StringToHGlobalAnsi(key);
-        string? result = await Command(args, 1, RequestType.Get);
-        _arrayPool.Return(args);
-        return result;
-    }
+        => await Command<string?>([key], RequestType.Get);
 
     private readonly object _lock = new();
 
     public void Dispose()
     {
+        GC.SuppressFinalize(this);
         lock (_lock)
         {
             if (_clientPointer == IntPtr.Zero)
@@ -79,14 +78,14 @@ public void Dispose()
     #endregion public methods
 
     #region private methods
-
+    // TODO rework the callback to handle other response types
     private void SuccessCallback(ulong index, IntPtr str)
     {
         string? result = str == IntPtr.Zero ? null : Marshal.PtrToStringAnsi(str);
         // Work needs to be offloaded from the calling thread, because otherwise we might starve the client's thread pool.
         _ = Task.Run(() =>
         {
-            Message<string> message = _messageContainer.GetMessage((int)index);
+            Message message = _messageContainer.GetMessage((int)index);
             message.SetResult(result);
         });
     }
@@ -95,7 +94,7 @@ private void FailureCallback(ulong index) =>
         // Work needs to be offloaded from the calling thread, because otherwise we might starve the client's thread pool.
         Task.Run(() =>
         {
-            Message<string> message = _messageContainer.GetMessage((int)index);
+            Message message = _messageContainer.GetMessage((int)index);
             message.SetException(new Exception("Operation failed"));
         });
 
@@ -114,7 +113,7 @@ private void FailureCallback(ulong index) =>
 
     /// Raw pointer to the underlying native client.
     private IntPtr _clientPointer;
-    private readonly MessageContainer<string> _messageContainer = new();
+    private readonly MessageContainer _messageContainer = new();
     private readonly ArrayPool<IntPtr> _arrayPool = ArrayPool<IntPtr>.Shared;
 
     #endregion private fields

csharp/lib/Commands/IGenericCommands.cs

@@ -0,0 +1,26 @@
+// Copyright Valkey GLIDE Project Contributors - SPDX Identifier: Apache-2.0
+
+namespace Glide.Commands;
+
+public interface IGenericCommands
+{
+    /// <summary>
+    /// Executes a single command, without checking inputs. Every part of the command, including subcommands,
+    /// should be added as a separate value in <paramref name="args"/>.
+    /// See the <see href="https://github.com/valkey-io/valkey-glide/wiki/General-Concepts#custom-command">Valkey GLIDE Wiki</see>.
+    /// for details on the restrictions and limitations of the custom command API.
+    /// <para />
+    /// This function should only be used for single-response commands. Commands that don't return complete response and awaits
+    /// (such as SUBSCRIBE), or that return potentially more than a single response (such as XREAD), or that change the client's
+    /// behavior (such as entering pub/sub mode on RESP2 connections) shouldn't be called using this function.
+    /// <example>
+    /// <code>
+    /// // Query all pub/sub clients
+    /// object result = await client.CustomCommand(["CLIENT", "LIST", "TYPE", "PUBSUB"]);
+    /// </code>
+    /// </example>
+    /// </summary>
+    /// <param name="args">A list including the command name and arguments for the custom command.</param>
+    /// <returns>The returning value depends on the executed command.</returns>
+    Task<object?> CustomCommand(string[] args);
+}

csharp/lib/Commands/IStringBaseCommands.cs

@@ -20,7 +20,7 @@ internal interface IStringBaseCommands
     /// <param name="key">The <paramref name="key"/> to store.</param>
     /// <param name="value">The value to store with the given <paramref name="key"/>.</param>
     /// <returns>A simple <c>"OK"</c> response.</returns>
-    Task<string?> Set(string key, string value);
+    Task<string> Set(string key, string value);
 
     /// <summary>
     /// Gets the value associated with the given <paramref name="key"/>, or <see langword="null"/> if no such <paramref name="key"/> exists.

csharp/lib/GlideClient.cs

@@ -6,6 +6,8 @@
 
 namespace Glide;
 
-public sealed class GlideClient(StandaloneClientConfiguration config) : BaseClient(config), IConnectionManagementCommands
+public sealed class GlideClient(StandaloneClientConfiguration config) : BaseClient(config), IConnectionManagementCommands, IGenericCommands
 {
+    public async Task<object?> CustomCommand(string[] args)
+        => await Command<object?>(args, RequestType.CustomCommand);
 }

csharp/lib/Internals/Message.cs

@@ -2,34 +2,35 @@
 
 using System.Diagnostics;
 using System.Runtime.CompilerServices;
-using System.Runtime.InteropServices;
 
 using Glide.Internals;
 
 /// Reusable source of ValueTask. This object can be allocated once and then reused
 /// to create multiple asynchronous operations, as long as each call to CreateTask
 /// is awaited to completion before the next call begins.
-internal class Message<T>(int index, MessageContainer<T> container) : INotifyCompletion
+internal class Message(int index, MessageContainer container) : INotifyCompletion
 {
     /// This is the index of the message in an external array, that allows the user to
     /// know how to find the message and set its result.
     public int Index { get; } = index;
-    /// The array holding the pointers  to the unmanaged memory that contains the operation's arguments.
+    /// The array holding the pointers to the unmanaged memory that contains the operation's arguments.
     public IntPtr[]? Args { get; private set; }
-    // We need to save the args count, because sometimes we get arrays that are larger than they need to be. We can't rely on `this.args.Length`, due to it coming from an array pool.
-    private int _argsCount;
-    private MessageContainer<T> Container { get; } = container;
+    private MessageContainer Container { get; } = container;
     private Action? _continuation = () => { };
     private const int COMPLETION_STAGE_STARTED = 0;
     private const int COMPLETION_STAGE_NEXT_SHOULD_EXECUTE_CONTINUATION = 1;
     private const int COMPLETION_STAGE_CONTINUATION_EXECUTED = 2;
     private int _completionState;
-    private T? _result;
+    private object? _result = default;
     private Exception? _exception;
+    // Holding the client prevents it from being GC'd until all operations complete.
+#pragma warning disable IDE0052 // Remove unread private members
+    private object? _client;
+#pragma warning restore IDE0052 // Remove unread private members
 
     /// Triggers a succesful completion of the task returned from the latest call
     /// to CreateTask.
-    public void SetResult(T? result)
+    public void SetResult(object? result)
     {
         _result = result;
         FinishSet();
@@ -45,7 +46,7 @@ public void SetException(Exception exc)
 
     private void FinishSet()
     {
-        FreePointers();
+        CleanUp();
 
         CheckRaceAndCallContinuation();
     }
@@ -67,44 +68,30 @@ private void CheckRaceAndCallContinuation()
         }
     }
 
-    public Message<T> GetAwaiter() => this;
+    public Message GetAwaiter() => this;
 
     /// This returns a task that will complete once SetException / SetResult are called,
     /// and ensures that the internal state of the message is set-up before the task is created,
     /// and cleaned once it is complete.
-    public void SetupTask(IntPtr[] arguments, int argsCount, object client)
+    public void SetupTask(IntPtr[] arguments, object client)
     {
         _continuation = null;
         _completionState = COMPLETION_STAGE_STARTED;
         _result = default;
         _exception = null;
         _client = client;
         Args = arguments;
-        _argsCount = argsCount;
     }
 
-    // This function isn't thread-safe. Access to it should be from a single thread, and only once per operation.
-    // For the sake of performance, this responsibility is on the caller, and the function doesn't contain any safety measures.
-    private void FreePointers()
+    private void CleanUp()
     {
         if (Args is { })
         {
-            for (int i = 0; i < _argsCount; i++)
-            {
-                Marshal.FreeHGlobal(Args[i]);
-            }
             Args = null;
-            _argsCount = 0;
         }
         _client = null;
     }
 
-    // Holding the client prevents it from being GC'd until all operations complete.
-#pragma warning disable IDE0052 // Remove unread private members
-    private object? _client;
-#pragma warning restore IDE0052 // Remove unread private members
-
-
     public void OnCompleted(Action continuation)
     {
         _continuation = continuation;
@@ -113,5 +100,5 @@ public void OnCompleted(Action continuation)
 
     public bool IsCompleted => _completionState == COMPLETION_STAGE_CONTINUATION_EXECUTED;
 
-    public T? GetResult() => _exception is null ? _result : throw _exception;
+    public object? GetResult() => _exception is null ? _result : throw _exception;
 }

csharp/lib/Internals/MessageContainer.cs

@@ -5,38 +5,39 @@
 namespace Glide.Internals;
 
 
-internal class MessageContainer<T>
+internal class MessageContainer
 {
-    internal Message<T> GetMessage(int index) => _messages[index];
+    internal Message GetMessage(int index) => _messages[index];
 
-    internal Message<T> GetMessageForCall(nint[] args, int argsCount)
+    internal Message GetMessageForCall<T>(nint[] args)
     {
-        Message<T> message = GetFreeMessage();
-        message.SetupTask(args, argsCount, this);
+        Message message = GetFreeMessage();
+        message.SetupTask(args, this);
         return message;
     }
 
-    private Message<T> GetFreeMessage()
+    private Message GetFreeMessage()
     {
-        if (!_availableMessages.TryDequeue(out Message<T>? message))
+        if (!_availableMessages.TryDequeue(out Message? message))
         {
             lock (_messages)
             {
                 int index = _messages.Count;
-                message = new Message<T>(index, this);
+                message = new Message(index, this);
                 _messages.Add(message);
             }
         }
         return message;
     }
 
-    public void ReturnFreeMessage(Message<T> message) => _availableMessages.Enqueue(message);
+    public void ReturnFreeMessage(Message message)
+        => _availableMessages.Enqueue((Message)(object)message);
 
     internal void DisposeWithError(Exception? error)
     {
         lock (_messages)
         {
-            foreach (Message<T>? message in _messages.Where(message => !message.IsCompleted))
+            foreach (Message? message in _messages.Where(message => !message.IsCompleted))
             {
                 try
                 {
@@ -52,9 +53,9 @@ internal void DisposeWithError(Exception? error)
     /// This list allows us random-access to the message in each index,
     /// which means that once we receive a callback with an index, we can
     /// find the message to resolve in constant time.
-    private readonly List<Message<T>> _messages = [];
+    private readonly List<Message> _messages = [];
 
     /// This queue contains the messages that were created and are currently unused by any task,
     /// so they can be reused y new tasks instead of allocating new messages.
-    private readonly ConcurrentQueue<Message<T>> _availableMessages = new();
+    private readonly ConcurrentQueue<Message> _availableMessages = new();
 }

csharp/lib/src/lib.rs

@@ -115,26 +115,20 @@ pub unsafe extern "C" fn command(
         Arc::from_raw(client_ptr as *mut Client)
     };
 
-    // The safety of these needs to be ensured by the calling code. Cannot dispose of the pointer before all operations have completed.
-    let args_address = args as usize;
+    let Some(mut cmd) = request_type.get_command() else {
+        unsafe {
+            (client.core.failure_callback)(callback_index); // TODO - report errors
+            return;
+        }
+    };
+    let args_slice = unsafe { std::slice::from_raw_parts(args, arg_count as usize) };
+    for arg in args_slice {
+        let c_str = unsafe { CStr::from_ptr(*arg as *mut c_char) };
+        cmd.arg(c_str.to_bytes());
+    }
 
     let core = client.core.clone();
     client.runtime.spawn(async move {
-        let Some(mut cmd) = request_type.get_command() else {
-            unsafe {
-                (core.failure_callback)(callback_index); // TODO - report errors
-                return;
-            }
-        };
-
-        let args_slice = unsafe {
-            std::slice::from_raw_parts(args_address as *const *mut c_char, arg_count as usize)
-        };
-        for arg in args_slice {
-            let c_str = unsafe { CStr::from_ptr(*arg as *mut c_char) };
-            cmd.arg(c_str.to_bytes());
-        }
-
         let result = core
             .client
             .clone()
@@ -145,7 +139,10 @@ pub unsafe extern "C" fn command(
             match result {
                 Ok(None) => (core.success_callback)(callback_index, std::ptr::null()),
                 Ok(Some(c_str)) => (core.success_callback)(callback_index, c_str.as_ptr()),
-                Err(_) => (core.failure_callback)(callback_index), // TODO - report errors
+                Err(err) => {
+                    dbg!(err); // TODO - report errors
+                    (core.failure_callback)(callback_index)
+                }
             };
         }
     });