diff --git a/Assets/Rivet/Editor/RivetToolchain.cs b/Assets/Rivet/Editor/RivetToolchain.cs index 95dbc52..e22181f 100644 --- a/Assets/Rivet/Editor/RivetToolchain.cs +++ b/Assets/Rivet/Editor/RivetToolchain.cs @@ -1,52 +1,159 @@ using System; +using System.Collections.Generic; using System.Diagnostics; using System.Runtime.InteropServices; +using Newtonsoft.Json.Linq; namespace Rivet.Editor { - public static class RivetToolchain + internal static class RivetToolchainFFI { const string RustLibrary = "__Internal"; - [UnmanagedFunctionPointer(CallingConvention.Cdecl)] - public delegate void EventCallback(ulong taskId, IntPtr eventJson); + [StructLayout(LayoutKind.Sequential)] + public struct RunTaskResult + { + public ulong TaskId; + public byte ErrorCode; + } [DllImport(RustLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "rivet_run_task")] - private static extern ulong run_task( + public static extern RunTaskResult run_task( [MarshalAs(UnmanagedType.LPStr)] string name, - [MarshalAs(UnmanagedType.LPStr)] string inputJson, - EventCallback callback + [MarshalAs(UnmanagedType.LPStr)] string inputJson ); [DllImport(RustLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "rivet_abort_task")] - private static extern bool abort_task(ulong taskId); + public static extern byte abort_task(ulong taskId); [DllImport(RustLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "rivet_shutdown")] - private static extern void shutdown(); + public static extern void shutdown(); [DllImport(RustLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "rivet_free_rust_string")] private static extern void free_rust_string(IntPtr str); - public static ulong RunTask(string name, string inputJson, EventCallback callback) + public static string PtrToString(IntPtr ptr) { - return run_task(name, inputJson, callback); + string result = Marshal.PtrToStringAnsi(ptr); + free_rust_string(ptr); + return result; } - public static bool AbortTask(ulong taskId) + [StructLayout(LayoutKind.Sequential)] + public struct TaskEvent { - return abort_task(taskId); + public ulong TaskId; + public IntPtr EventJson; + } + + [StructLayout(LayoutKind.Sequential)] + public struct PollTaskEventsResult + { + public UIntPtr Count; + public byte ErrorCode; + } + + [DllImport(RustLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "rivet_poll_task_events")] + public static extern PollTaskEventsResult poll_task_events( + [In, Out] TaskEvent[] events, + UIntPtr maxEvents + ); + + public const uint POLL_TASK_EVENT_BUFFER = 128; + } + + public static class RivetToolchain + { + private static Dictionary> taskCallbacks = new Dictionary>(); + private static readonly object taskCallbacksLock = new object(); + + public static ulong RunTask(string name, string inputJson, Action callback) + { + RivetToolchainFFI.RunTaskResult result = RivetToolchainFFI.run_task(name, inputJson); + if (result.ErrorCode != 0) + { + throw new Exception($"Rivet task '{name}' failed with error code: {result.ErrorCode}"); + } + + lock (taskCallbacksLock) + { + taskCallbacks[result.TaskId] = callback; + } + + return result.TaskId; + } + + public static void AbortTask(ulong taskId) + { + var error_code = RivetToolchainFFI.abort_task(taskId); + if (error_code != 0) + { + throw new Exception($"Failed to abort Rivet task with ID: {taskId}"); + } } public static void Shutdown() { - // shutdown(); + // RivetToolchainFFI.shutdown(); } - public static string PtrToString(IntPtr ptr) + public class TaskEventData { - string result = Marshal.PtrToStringAnsi(ptr); - free_rust_string(ptr); - return result; + public ulong TaskId; + public JObject EventJson; + } + + public static List PollTaskEvents() + { + List allEvents = new List(); + RivetToolchainFFI.TaskEvent[] buffer = new RivetToolchainFFI.TaskEvent[RivetToolchainFFI.POLL_TASK_EVENT_BUFFER]; + + uint totalPolled = 0; + while (true) + { + RivetToolchainFFI.PollTaskEventsResult result = RivetToolchainFFI.poll_task_events(buffer, new UIntPtr(RivetToolchainFFI.POLL_TASK_EVENT_BUFFER)); + + if (result.ErrorCode != 0) + { + throw new Exception($"Failed to poll Rivet task events with error code: {result.ErrorCode}"); + } + + uint polledCount = result.Count.ToUInt32(); + totalPolled += polledCount; + + for (int i = 0; i < polledCount; i++) + { + string eventJsonString = RivetToolchainFFI.PtrToString(buffer[i].EventJson); + JObject eventJson = JObject.Parse(eventJsonString); + + // Attempt to get callback + Action callback = null; + lock (taskCallbacksLock) + { + if (!taskCallbacks.TryGetValue(buffer[i].TaskId, out callback)) + { + RivetLogger.Warning($"Missing callback for task ID: {buffer[i].TaskId}"); + } + } + + // Call callback + callback?.Invoke(eventJson); + } + + if (polledCount == 0) + { + // No more events to poll + break; + } + } + + // Log task event + if (totalPolled > 0) + { + RivetLogger.Log($"Polled {totalPolled} task events"); + } + + return allEvents; } } } \ No newline at end of file diff --git a/Assets/Rivet/Editor/Task.cs b/Assets/Rivet/Editor/Task.cs index 22599fa..cdd756c 100644 --- a/Assets/Rivet/Editor/Task.cs +++ b/Assets/Rivet/Editor/Task.cs @@ -49,15 +49,8 @@ private void Run(string name, string inputJson) _taskId = RivetToolchain.RunTask(name, inputJson, OnOutputEvent); } - private void OnOutputEvent(ulong taskId, IntPtr eventJsonPtr) + private void OnOutputEvent(JObject eventObj) { - string eventJson = RivetToolchain.PtrToString(eventJsonPtr); - EditorApplication.delayCall += () => HandleOnOutputEvent(eventJson); - } - - private void HandleOnOutputEvent(string eventJson) - { - var eventObj = JObject.Parse(eventJson); if (eventObj.ContainsKey("log")) { OnLogEvent((string)eventObj["log"]); @@ -85,7 +78,7 @@ private void HandleOnOutputEvent(string eventJson) } else { - RivetLogger.Warning($"Unknown event {eventJson}"); + RivetLogger.Warning($"Unknown event {eventObj.ToString(Newtonsoft.Json.Formatting.None)}"); } } diff --git a/Assets/Rivet/Editor/UI/Dock/Dock.cs b/Assets/Rivet/Editor/UI/Dock/Dock.cs index dc45c87..ec68bc6 100644 --- a/Assets/Rivet/Editor/UI/Dock/Dock.cs +++ b/Assets/Rivet/Editor/UI/Dock/Dock.cs @@ -181,6 +181,9 @@ public void OnEnable() autoRestart: true ); + // Add update callback + EditorApplication.update += OnEditorUpdate; + // Shut down on reload AssemblyReloadEvents.beforeAssemblyReload += () => { @@ -199,6 +202,9 @@ public void OnEnable() public void OnDisable() { RivetLogger.Log("On Disable"); + + EditorApplication.update -= OnEditorUpdate; + ShutdownPlugin(); } @@ -305,5 +311,10 @@ private void OpenHub() Application.OpenURL($"https://hub.rivet.gg"); } } + + private void OnEditorUpdate() + { + RivetToolchain.PollTaskEvents(); + } } }