Make tokenizer per-thread to unblock concurrent Redact

src/PrivacyFilter.Net/PrivacyFilter.cs

@@ -132,6 +132,7 @@ public void Dispose()
         }
 
         _session.Dispose();
+        _tokenizer.Dispose();
         _disposed = true;
     }
 

src/PrivacyFilter.Net/PrivacyFilterTokenizer.cs

@@ -2,61 +2,61 @@
 
 namespace PrivacyFilterNet;
 
-internal sealed class PrivacyFilterTokenizer
+internal sealed class PrivacyFilterTokenizer : IDisposable
 {
-    private readonly TiktokenTokenizer _tokenizer;
-    private readonly object _gate = new();
-
-    public PrivacyFilterTokenizer()
-    {
-        _tokenizer = TiktokenTokenizer.CreateForEncoding("o200k_base");
-    }
+    // TiktokenTokenizer keeps internal encode caches that are not guaranteed safe
+    // for concurrent mutation, so give each thread its own instance rather than
+    // serializing all tokenization behind a global lock. Combined with ONNX
+    // Runtime's thread-safe Run, this lets a single shared PrivacyFilter serve many
+    // concurrent Redact calls at full throughput instead of bottlenecking here.
+    private readonly ThreadLocal<TiktokenTokenizer> _tokenizer =
+        new(() => TiktokenTokenizer.CreateForEncoding("o200k_base"));
 
     public TokenizedText Encode(string text)
     {
         ArgumentNullException.ThrowIfNull(text);
 
-        lock (_gate)
+        TiktokenTokenizer tokenizer = _tokenizer.Value!;
+        IReadOnlyList<EncodedToken> encodedTokens =
+            tokenizer.EncodeToTokens(text, out _, considerPreTokenization: true, considerNormalization: true);
+        int[] ids = encodedTokens.Select(token => token.Id).ToArray();
+        if (ids.Length == 0)
         {
-            IReadOnlyList<EncodedToken> encodedTokens =
-                _tokenizer.EncodeToTokens(text, out _, considerPreTokenization: true, considerNormalization: true);
-            int[] ids = encodedTokens.Select(token => token.Id).ToArray();
-            if (ids.Length == 0)
-            {
-                return new TokenizedText(ids, text, [], [], DecodedMismatch: false);
-            }
-
-            string decodedText = _tokenizer.Decode(ids);
-            bool decodedMismatch = !string.Equals(decodedText, text, StringComparison.Ordinal);
-            if (decodedMismatch)
-            {
-                encodedTokens =
-                    _tokenizer.EncodeToTokens(decodedText, out _, considerPreTokenization: true, considerNormalization: true);
-                if (encodedTokens.Count != ids.Length ||
-                    !encodedTokens.Select(token => token.Id).SequenceEqual(ids))
-                {
-                    throw new InvalidDataException(
-                        "Tokenizer decode did not produce a stable token sequence for span offsets.");
-                }
-            }
+            return new TokenizedText(ids, text, [], [], DecodedMismatch: false);
+        }
 
-            var charStarts = new int[encodedTokens.Count];
-            var charEnds = new int[encodedTokens.Count];
-            for (int index = 0; index < encodedTokens.Count; index++)
+        string decodedText = tokenizer.Decode(ids);
+        bool decodedMismatch = !string.Equals(decodedText, text, StringComparison.Ordinal);
+        if (decodedMismatch)
+        {
+            encodedTokens =
+                tokenizer.EncodeToTokens(decodedText, out _, considerPreTokenization: true, considerNormalization: true);
+            if (encodedTokens.Count != ids.Length ||
+                !encodedTokens.Select(token => token.Id).SequenceEqual(ids))
             {
-                Range offset = encodedTokens[index].Offset;
-                charStarts[index] = offset.Start.GetOffset(decodedText.Length);
-                charEnds[index] = offset.End.GetOffset(decodedText.Length);
+                throw new InvalidDataException(
+                    "Tokenizer decode did not produce a stable token sequence for span offsets.");
             }
+        }
 
-            return new TokenizedText(
-                ids,
-                decodedText,
-                charStarts,
-                charEnds,
-                DecodedMismatch: decodedMismatch);
+        var charStarts = new int[encodedTokens.Count];
+        var charEnds = new int[encodedTokens.Count];
+        for (int index = 0; index < encodedTokens.Count; index++)
+        {
+            Range offset = encodedTokens[index].Offset;
+            charStarts[index] = offset.Start.GetOffset(decodedText.Length);
+            charEnds[index] = offset.End.GetOffset(decodedText.Length);
         }
+
+        return new TokenizedText(
+            ids,
+            decodedText,
+            charStarts,
+            charEnds,
+            DecodedMismatch: decodedMismatch);
     }
+
+    public void Dispose() => _tokenizer.Dispose();
 }
 
 internal sealed record TokenizedText(

tests/PrivacyFilter.Net.Tests/TokenizerTests.cs

@@ -25,6 +25,40 @@ public void MatchesTiktokenOracle()
         }
     }
 
+    [Fact]
+    public void EncodeIsThreadSafeUnderConcurrency()
+    {
+        // The tokenizer is now per-thread (replacing a global lock), so a single
+        // shared instance must serve concurrent Encode calls without corruption.
+        // Hammer it from many threads against the oracle: mismatched ids/offsets or
+        // an exception would reveal unsafe shared mutable state.
+        string path = Path.Combine(AppContext.BaseDirectory, "tokenizer_oracle.json");
+        TokenizerOracleCase[] cases =
+            JsonSerializer.Deserialize<TokenizerOracleCase[]>(
+                File.ReadAllBytes(path),
+                new JsonSerializerOptions { PropertyNameCaseInsensitive = true })!;
+        using var tokenizer = new PrivacyFilterTokenizer();
+
+        Parallel.For(
+            0,
+            256,
+            new ParallelOptions
+            {
+                MaxDegreeOfParallelism = Math.Max(4, Environment.ProcessorCount),
+            },
+            _ =>
+            {
+                foreach (TokenizerOracleCase testCase in cases)
+                {
+                    TokenizedText actual = tokenizer.Encode(testCase.Text);
+                    Assert.Equal(testCase.Ids, actual.TokenIds);
+                    Assert.Equal(testCase.Decoded, actual.DecodedText);
+                    Assert.Equal(testCase.Starts, actual.CharacterStarts);
+                    Assert.Equal(testCase.Ends, actual.CharacterEnds);
+                }
+            });
+    }
+
     private sealed record TokenizerOracleCase(
         string Text,
         int[] Ids,