1 files changed, 7 insertions, 6 deletions
diff --git a/src/csharp/Grpc.Microbenchmarks/CompletionRegistryBenchmark.cs b/src/csharp/Grpc.Microbenchmarks/CompletionRegistryBenchmark.cs
index 0e2fb070f1..2d1c33e9a0 100644
--- a/src/csharp/Grpc.Microbenchmarks/CompletionRegistryBenchmark.cs
+++ b/src/csharp/Grpc.Microbenchmarks/CompletionRegistryBenchmark.cs
@@ -40,24 +40,25 @@ namespace Grpc.Microbenchmarks
             GrpcEnvironment.ReleaseAsync().Wait();
         }
 
-        public void Run(int threadCount, int iterations)
+        public void Run(int threadCount, int iterations, bool useSharedRegistry)
         {
-            Console.WriteLine(string.Format("CompletionRegistryBenchmark: threads={0}, iterations={1}", threadCount, iterations));
-            var threadedBenchmark = new ThreadedBenchmark(threadCount, () => ThreadBody(iterations));
+            Console.WriteLine(string.Format("CompletionRegistryBenchmark: threads={0}, iterations={1}, useSharedRegistry={2}", threadCount, iterations, useSharedRegistry));
+            CompletionRegistry sharedRegistry = useSharedRegistry ? new CompletionRegistry(environment) : null;
+            var threadedBenchmark = new ThreadedBenchmark(threadCount, () => ThreadBody(iterations, sharedRegistry));
             threadedBenchmark.Run();
             // TODO: parametrize by number of pending completions
         }
 
-        private void ThreadBody(int iterations)
+        private void ThreadBody(int iterations, CompletionRegistry optionalSharedRegistry)
         {
-            var completionRegistry = new CompletionRegistry(environment);
+            var completionRegistry = optionalSharedRegistry ?? new CompletionRegistry(environment);
             var ctx = BatchContextSafeHandle.Create();
   
             var stopwatch = Stopwatch.StartNew();
             for (int i = 0; i < iterations; i++)
             {
                 completionRegistry.Register(ctx.Handle, ctx);
-                var callback = completionRegistry.Extract(completionRegistry.LastRegisteredKey);
+                var callback = completionRegistry.Extract(ctx.Handle);
                 // NOTE: we are not calling the callback to avoid disposing ctx.
             }
             stopwatch.Stop();