Callbacks working / test passes

Author: Sushisource

.gitignore

@@ -7,4 +7,6 @@ obj/
 /tests/golangworker/golangworker
 /.vs
 /.vscode
-/.idea
+/.idea
+/.zed
+Temporalio.sln.DotSettings.user

src/Temporalio/Bridge/CustomSlotSupplier.cs

@@ -1,3 +1,6 @@
+using System;
+using System.Collections.Generic;
+using System.Runtime.InteropServices;
 using System.Threading.Tasks;
 
 namespace Temporalio.Bridge
@@ -8,6 +11,8 @@ namespace Temporalio.Bridge
     internal class CustomSlotSupplier : NativeInvokeableClass<Interop.CustomSlotSupplierCallbacks>
     {
         private readonly Temporalio.Worker.Tuning.ICustomSlotSupplier userSupplier;
+        private readonly Dictionary<uint, GCHandle> permits = new();
+        private uint permitId = 1;
 
         /// <summary>
         /// Initializes a new instance of the <see cref="CustomSlotSupplier" /> class.
@@ -28,25 +33,65 @@ internal unsafe CustomSlotSupplier(Temporalio.Worker.Tuning.ICustomSlotSupplier
             PinCallbackHolder(interopCallbacks);
         }
 
-        private void Reserve(Interop.SlotReserveCtx ctx)
+        private unsafe void Reserve(Interop.SlotReserveCtx ctx, void* sender)
         {
-            // TODO: Need to call callback with result that will put it in a channel to await in Rust
-            var reserveTask = Task.Run(() => userSupplier.ReserveSlotAsync(new(ctx)));
+            SafeReserve(ctx, new IntPtr(sender));
         }
 
-        private void TryReserve(Interop.SlotReserveCtx ctx)
+        private void SafeReserve(Interop.SlotReserveCtx ctx, IntPtr sender)
         {
-            userSupplier.TryReserveSlot(new(ctx));
+            var reserveTask = Task.Run(async () =>
+            {
+                try
+                {
+                    var permit = await userSupplier.ReserveSlotAsync(new(ctx)).ConfigureAwait(false);
+                    var usedPermitId = AddPermitToMap(permit);
+                    unsafe
+                    {
+                        Interop.Methods.complete_async_reserve(sender.ToPointer(), new(usedPermitId));
+                    }
+                }
+                catch (Exception e)
+                {
+                    Console.WriteLine("Exception in reserve: " + e.Message);
+                    throw;
+                }
+            });
+        }
+
+        private unsafe UIntPtr TryReserve(Interop.SlotReserveCtx ctx)
+        {
+            var maybePermit = userSupplier.TryReserveSlot(new(ctx));
+            if (maybePermit == null)
+            {
+                return UIntPtr.Zero;
+            }
+            var usedPermitId = AddPermitToMap(maybePermit);
+            return new(usedPermitId);
         }
 
         private void MarkUsed(Interop.SlotMarkUsedCtx ctx)
         {
-            userSupplier.MarkSlotUsed(new(ctx));
+            userSupplier.MarkSlotUsed(new(ctx, permits[ctx.slot_permit.ToUInt32()]));
         }
 
         private void Release(Interop.SlotReleaseCtx ctx)
         {
-            userSupplier.ReleaseSlot(new(ctx));
+            var permitId = ctx.slot_permit.ToUInt32();
+            userSupplier.ReleaseSlot(new(ctx, permits[permitId]));
+            permits.Remove(permitId);
+        }
+
+        private uint AddPermitToMap(Temporalio.Worker.Tuning.SlotPermit permit)
+        {
+            var handle = GCHandle.Alloc(permit);
+            lock (permits)
+            {
+                var usedPermitId = permitId;
+                permits.Add(permitId, handle);
+                permitId += 1;
+                return usedPermitId;
+            }
         }
     }
 }

src/Temporalio/Bridge/Interop/Interop.cs

@@ -577,10 +577,11 @@ internal partial struct SlotReserveCtx
     }
 
     [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
-    internal delegate void CustomReserveSlotCallback([NativeTypeName("struct SlotReserveCtx")] SlotReserveCtx ctx);
+    internal unsafe delegate void CustomReserveSlotCallback([NativeTypeName("struct SlotReserveCtx")] SlotReserveCtx ctx, void* sender);
 
     [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
-    internal delegate void CustomTryReserveSlotCallback([NativeTypeName("struct SlotReserveCtx")] SlotReserveCtx ctx);
+    [return: NativeTypeName("uintptr_t")]
+    internal delegate UIntPtr CustomTryReserveSlotCallback([NativeTypeName("struct SlotReserveCtx")] SlotReserveCtx ctx);
 
     internal enum SlotInfo_Tag
     {
@@ -614,7 +615,7 @@ internal unsafe partial struct SlotInfo
     {
         public SlotInfo_Tag tag;
 
-        [NativeTypeName("__AnonymousRecord_temporal-sdk-bridge_L419_C3")]
+        [NativeTypeName("__AnonymousRecord_temporal-sdk-bridge_L422_C3")]
         public _Anonymous_e__Union Anonymous;
 
         internal ref WorkflowSlotInfo_Body workflow_slot_info
@@ -664,13 +665,13 @@ internal partial struct _Anonymous_e__Union
         }
     }
 
-    internal unsafe partial struct SlotMarkUsedCtx
+    internal partial struct SlotMarkUsedCtx
     {
         [NativeTypeName("struct SlotInfo")]
         public SlotInfo slot_info;
 
-        [NativeTypeName("const void *")]
-        public void* slot_permit;
+        [NativeTypeName("uintptr_t")]
+        public UIntPtr slot_permit;
     }
 
     [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
@@ -681,8 +682,8 @@ internal unsafe partial struct SlotReleaseCtx
         [NativeTypeName("const struct SlotInfo *")]
         public SlotInfo* slot_info;
 
-        [NativeTypeName("const void *")]
-        public void* slot_permit;
+        [NativeTypeName("uintptr_t")]
+        public UIntPtr slot_permit;
     }
 
     [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
@@ -720,7 +721,7 @@ internal unsafe partial struct SlotSupplier
     {
         public SlotSupplier_Tag tag;
 
-        [NativeTypeName("__AnonymousRecord_temporal-sdk-bridge_L465_C3")]
+        [NativeTypeName("__AnonymousRecord_temporal-sdk-bridge_L468_C3")]
         public _Anonymous_e__Union Anonymous;
 
         internal ref FixedSizeSlotSupplier fixed_size
@@ -760,15 +761,15 @@ internal ref CustomSlotSupplierCallbacksImpl custom
         internal unsafe partial struct _Anonymous_e__Union
         {
             [FieldOffset(0)]
-            [NativeTypeName("__AnonymousRecord_temporal-sdk-bridge_L466_C5")]
+            [NativeTypeName("__AnonymousRecord_temporal-sdk-bridge_L469_C5")]
             public _Anonymous1_e__Struct Anonymous1;
 
             [FieldOffset(0)]
-            [NativeTypeName("__AnonymousRecord_temporal-sdk-bridge_L469_C5")]
+            [NativeTypeName("__AnonymousRecord_temporal-sdk-bridge_L472_C5")]
             public _Anonymous2_e__Struct Anonymous2;
 
             [FieldOffset(0)]
-            [NativeTypeName("__AnonymousRecord_temporal-sdk-bridge_L472_C5")]
+            [NativeTypeName("__AnonymousRecord_temporal-sdk-bridge_L475_C5")]
             public _Anonymous3_e__Struct Anonymous3;
 
             internal partial struct _Anonymous1_e__Struct
@@ -1057,5 +1058,8 @@ internal static unsafe partial class Methods
         [DllImport("temporal_sdk_bridge", CallingConvention = CallingConvention.Cdecl, ExactSpelling = true)]
         [return: NativeTypeName("struct WorkerReplayPushResult")]
         public static extern WorkerReplayPushResult worker_replay_push([NativeTypeName("struct Worker *")] Worker* worker, [NativeTypeName("struct WorkerReplayPusher *")] WorkerReplayPusher* worker_replay_pusher, [NativeTypeName("struct ByteArrayRef")] ByteArrayRef workflow_id, [NativeTypeName("struct ByteArrayRef")] ByteArrayRef history);
+
+        [DllImport("temporal_sdk_bridge", CallingConvention = CallingConvention.Cdecl, ExactSpelling = true)]
+        public static extern void complete_async_reserve(void* sender, [NativeTypeName("uintptr_t")] UIntPtr permit_id);
     }
 }

src/Temporalio/Bridge/include/temporal-sdk-bridge.h

@@ -391,9 +391,12 @@ typedef struct SlotReserveCtx {
   bool is_sticky;
 } SlotReserveCtx;
 
-typedef void (*CustomReserveSlotCallback)(struct SlotReserveCtx ctx);
+typedef void (*CustomReserveSlotCallback)(struct SlotReserveCtx ctx, void *sender);
 
-typedef void (*CustomTryReserveSlotCallback)(struct SlotReserveCtx ctx);
+/**
+ * Must return C#-tracked id for the permit. A zero value means no permit was reserved.
+ */
+typedef uintptr_t (*CustomTryReserveSlotCallback)(struct SlotReserveCtx ctx);
 
 typedef enum SlotInfo_Tag {
   WorkflowSlotInfo,
@@ -426,19 +429,19 @@ typedef struct SlotInfo {
 typedef struct SlotMarkUsedCtx {
   struct SlotInfo slot_info;
   /**
-   * User instance of a slot permit.
+   * C# id for the slot permit.
    */
-  const void *slot_permit;
+  uintptr_t slot_permit;
 } SlotMarkUsedCtx;
 
 typedef void (*CustomMarkSlotUsedCallback)(struct SlotMarkUsedCtx ctx);
 
 typedef struct SlotReleaseCtx {
   const struct SlotInfo *slot_info;
   /**
-   * User instance of a slot permit.
+   * C# id for the slot permit.
    */
-  const void *slot_permit;
+  uintptr_t slot_permit;
 } SlotReleaseCtx;
 
 typedef void (*CustomReleaseSlotCallback)(struct SlotReleaseCtx ctx);
@@ -689,6 +692,8 @@ struct WorkerReplayPushResult worker_replay_push(struct Worker *worker,
                                                  struct ByteArrayRef workflow_id,
                                                  struct ByteArrayRef history);
 
+void complete_async_reserve(void *sender, uintptr_t permit_id);
+
 #ifdef __cplusplus
 } // extern "C"
 #endif // __cplusplus

src/Temporalio/Bridge/src/lib.rs

@@ -31,13 +31,6 @@ impl ByteArrayRef {
         }
     }
 
-    fn from_string(s: &String) -> ByteArrayRef {
-        ByteArrayRef {
-            data: s.as_ptr(),
-            size: s.len(),
-        }
-    }
-
     fn to_slice(&self) -> &[u8] {
         unsafe { std::slice::from_raw_parts(self.data, self.size) }
     }
@@ -113,7 +106,7 @@ impl From<&str> for ByteArrayRef {
 
 impl From<String> for ByteArrayRef {
     fn from(s: String) -> ByteArrayRef {
-        ByteArrayRef::from_string(&s)
+        ByteArrayRef::from_str(&s)
     }
 }
 

src/Temporalio/Bridge/src/metric.rs

@@ -417,7 +417,7 @@ impl CustomMetricMeterRef {
                         ),
                     };
                     CustomMetricAttribute {
-                        key: ByteArrayRef::from_string(&kv.key),
+                        key: ByteArrayRef::from_str(&kv.key),
                         value,
                         value_type,
                     }

src/Temporalio/Bridge/src/runtime.rs

@@ -316,13 +316,13 @@ impl fmt::Debug for LogForwarder {
 #[no_mangle]
 pub extern "C" fn forwarded_log_target(log: *const ForwardedLog) -> ByteArrayRef {
     let log = unsafe { &*log };
-    ByteArrayRef::from_string(&log.core.target)
+    ByteArrayRef::from_str(&log.core.target)
 }
 
 #[no_mangle]
 pub extern "C" fn forwarded_log_message(log: *const ForwardedLog) -> ByteArrayRef {
     let log = unsafe { &*log };
-    ByteArrayRef::from_string(&log.core.message)
+    ByteArrayRef::from_str(&log.core.message)
 }
 
 #[no_mangle]