perf(replay): encode screenshots and batches off the main thread

.changeset/replay-capture-off-main-thread.md

@@ -0,0 +1,5 @@
+---
+'com.posthog.unity': patch
+---
+
+Reduce session replay frame hitches by encoding screenshots to JPEG and base64 on a worker thread, and by serializing and gzipping replay batches on a worker thread instead of the main thread.

com.posthog.unity/Runtime/SessionReplay/ReplayQueue.cs

@@ -4,6 +4,7 @@
 using System.IO;
 using System.IO.Compression;
 using System.Text;
+using System.Threading.Tasks;
 using UnityEngine;
 using UnityEngine.Networking;
 
@@ -20,7 +21,7 @@ class ReplayQueue
         readonly string _host;
         readonly Func<string> _getDistinctId;
         readonly Func<string> _getSessionId;
-        readonly Func<List<SnapshotEvent>, Action<bool, int>, IEnumerator> _sendBatch;
+        readonly Func<List<SnapshotEvent>, bool, Action<bool, int>, IEnumerator> _sendBatch;
         readonly Func<DateTime> _utcNow;
         readonly Func<bool> _isConnected;
         readonly object _lock = new();
@@ -62,7 +63,7 @@ internal ReplayQueue(
             string host,
             Func<string> getDistinctId,
             Func<string> getSessionId,
-            Func<List<SnapshotEvent>, Action<bool, int>, IEnumerator> sendBatch,
+            Func<List<SnapshotEvent>, bool, Action<bool, int>, IEnumerator> sendBatch,
             Func<DateTime> utcNow,
             Func<bool> isConnected
         )
@@ -156,7 +157,9 @@ public int Count
         /// <summary>
         /// Flushes all events in the queue.
         /// </summary>
-        public void Flush()
+        /// <param name="prepareInline">Serialize and compress the first batch on the main thread, so its request
+        /// starts before the coroutine first yields. Use it when the app is about to be suspended.</param>
+        public void Flush(bool prepareInline = false)
         {
             if (!_isRunning)
             {
@@ -166,7 +169,7 @@ public void Flush()
 
             if (_coroutineRunner != null)
             {
-                _coroutineRunner.StartCoroutine(FlushCoroutine());
+                _coroutineRunner.StartCoroutine(FlushCoroutine(prepareInline));
             }
         }
 
@@ -226,10 +229,14 @@ IEnumerator FlushTimerCoroutine()
             }
         }
 
-        internal IEnumerator FlushCoroutine()
+        internal IEnumerator FlushCoroutine(bool prepareInline = false)
         {
             lock (_lock)
             {
+                // A flush already in flight makes an inline flush skip too. If it is still preparing on the
+                // worker when the app is suspended, its batch waits until the app resumes. That window is
+                // the prepare time (~100 ms) of each flush, so it is accepted rather than finished
+                // synchronously. A flush waiting on its response skips this one the same way main did.
                 if (_isFlushing)
                 {
                     PostHogLogger.Debug("Already flushing replay queue, skipping");
@@ -276,6 +283,7 @@ internal IEnumerator FlushCoroutine()
                     var statusCode = 0;
                     yield return _sendBatch(
                         batch,
+                        prepareInline,
                         (result, code) =>
                         {
                             success = result;
@@ -288,6 +296,9 @@ internal IEnumerator FlushCoroutine()
                         RemoveBatch(batch);
                         ResetRetryState();
                         PostHogLogger.Debug($"Successfully sent {batch.Count} replay events");
+                        // Only the first batch has to start before a suspend, and later batches would block
+                        // the main thread again.
+                        prepareInline = false;
                         continue;
                     }
 
@@ -356,7 +367,11 @@ void ResetRetryState()
             _pausedUntil = null;
         }
 
-        IEnumerator SendBatch(List<SnapshotEvent> events, Action<bool, int> onComplete)
+        IEnumerator SendBatch(
+            List<SnapshotEvent> events,
+            bool prepareInline,
+            Action<bool, int> onComplete
+        )
         {
             var url = $"{_host}/s/";
 
@@ -366,10 +381,42 @@ IEnumerator SendBatch(List<SnapshotEvent> events, Action<bool, int> onComplete)
                 batchList.Add(evt.ToDictionary(_apiKey));
             }
 
-            var json = JsonSerializer.Serialize(batchList);
-            var bodyBytes = Encoding.UTF8.GetBytes(json);
+            Task<(byte[] PayloadBytes, bool UseCompression, Exception CompressionError)> prepare;
+            if (prepareInline)
+            {
+                // The app is about to be suspended, and Unity's loop stops with it. Preparing here lets the
+                // request start before the first yield, so the OS sends it before suspending the app.
+                prepare = PrepareBodyInline(batchList);
+            }
+            else
+            {
+                // Serializing and compressing a batch of screenshots takes over 100 ms on mobile, so it runs
+                // on a worker thread while this coroutine waits. None of it touches Unity APIs.
+                prepare = Task.Run(() => PrepareBody(batchList));
+                while (!prepare.IsCompleted)
+                {
+                    yield return null;
+                }
+            }
 
-            var (payloadBytes, useCompression) = PreparePayload(bodyBytes, CompressGzip);
+            if (prepare.IsFaulted || prepare.IsCanceled)
+            {
+                PostHogLogger.Error(
+                    "Failed to serialize replay batch",
+                    prepare.Exception?.GetBaseException()
+                );
+                // The same batch would fail again, so drop it like a rejected request instead of retrying it.
+                onComplete?.Invoke(false, 400);
+                yield break;
+            }
+
+            var (payloadBytes, useCompression, compressionError) = prepare.Result;
+            if (compressionError != null)
+            {
+                PostHogLogger.Warning(
+                    $"Failed to gzip replay batch, sending uncompressed: {compressionError.Message}"
+                );
+            }
             PostHogLogger.Debug(
                 $"Sending replay batch to {url} (size: {payloadBytes.Length} bytes)"
             );
@@ -405,25 +452,57 @@ IEnumerator SendBatch(List<SnapshotEvent> events, Action<bool, int> onComplete)
             }
         }
 
-        internal static (byte[] PayloadBytes, bool UseCompression) PreparePayload(
-            byte[] bodyBytes,
-            Func<byte[], byte[]> compressGzip
-        )
+        /// <summary>
+        /// Runs <see cref="PrepareBody"/> on the calling thread, returning a failure as a faulted task so the
+        /// caller handles both paths the same way.
+        /// </summary>
+        internal static Task<(
+            byte[] PayloadBytes,
+            bool UseCompression,
+            Exception CompressionError
+        )> PrepareBodyInline(List<Dictionary<string, object>> batchList)
         {
             try
             {
-                return (compressGzip(bodyBytes), true);
+                return Task.FromResult(PrepareBody(batchList));
             }
             catch (Exception ex)
             {
-                PostHogLogger.Warning(
-                    $"Failed to gzip replay batch, sending uncompressed: {ex.Message}"
-                );
-                return (bodyBytes, false);
+                return Task.FromException<(byte[], bool, Exception)>(ex);
+            }
+        }
+
+        /// <summary>
+        /// Runs on a worker thread, so it doesn't log; the caller reports a compression error.
+        /// </summary>
+        internal static (
+            byte[] PayloadBytes,
+            bool UseCompression,
+            Exception CompressionError
+        ) PrepareBody(List<Dictionary<string, object>> batchList)
+        {
+            var json = JsonSerializer.Serialize(batchList);
+            var bodyBytes = Encoding.UTF8.GetBytes(json);
+            return PreparePayload(bodyBytes, CompressGzip);
+        }
+
+        internal static (
+            byte[] PayloadBytes,
+            bool UseCompression,
+            Exception CompressionError
+        ) PreparePayload(byte[] bodyBytes, Func<byte[], byte[]> compressGzip)
+        {
+            try
+            {
+                return (compressGzip(bodyBytes), true, null);
+            }
+            catch (Exception ex)
+            {
+                return (bodyBytes, false, ex);
             }
         }
 
-        byte[] CompressGzip(byte[] data)
+        static byte[] CompressGzip(byte[] data)
         {
             using var output = new MemoryStream();
             using (var gzip = new GZipStream(output, CompressionMode.Compress))