Fix race condition in SingleThreadSynchronizationContext (#4468)

vonzshik · vonzshik · commit 3af871d5a017 · 2022-06-19T15:41:10.000+03:00
Fixes #4465 (cherry picked from commit 4e7f2a7)
diff --git a/src/Npgsql/SingleThreadSynchronizationContext.cs b/src/Npgsql/SingleThreadSynchronizationContext.cs
@@ -1,5 +1,6 @@
 ﻿using System;
 using System.Collections.Concurrent;
+using System.Diagnostics;
 using System.Threading;
 using JetBrains.Annotations;
 
@@ -8,8 +9,10 @@ namespace Npgsql
     sealed class SingleThreadSynchronizationContext : SynchronizationContext, IDisposable
     {
         readonly BlockingCollection<CallbackAndState> _tasks = new BlockingCollection<CallbackAndState>();
+        readonly object _lockObject = new object();
         [CanBeNull]
-        Thread _thread;
+        volatile Thread _thread;
+        bool _doingWork;
 
         const int ThreadStayAliveMs = 10000;
         readonly string _threadName;
@@ -23,12 +26,16 @@ public override void Post(SendOrPostCallback callback, object state)
         {
             _tasks.Add(new CallbackAndState { Callback = callback, State = state });
 
-            if (_thread == null)
+            lock (_lockObject)
             {
-                lock (this)
+                if (!_doingWork)
                 {
-                    if (_thread != null)
-                        return;
+                    // Either there is no thread, or the current thread is exiting
+                    // In which case, wait for it to complete
+                    var currentThread = _thread;
+                    currentThread?.Join();
+                    Debug.Assert(_thread is null);
+                    _doingWork = true;
                     _thread = new Thread(WorkLoop) { Name = _threadName, IsBackground = true };
                     _thread.Start();
                 }
@@ -38,10 +45,10 @@ public override void Post(SendOrPostCallback callback, object state)
         public void Dispose()
         {
             _tasks.CompleteAdding();
-            lock (this)
-            {
-                _thread?.Join();
-            }
+            var thread = _thread;
+            thread?.Join();
+
+            _tasks.Dispose();
         }
 
         void WorkLoop()
@@ -54,13 +61,40 @@ void WorkLoop()
                 {
                     var taken = _tasks.TryTake(out var callbackAndState, ThreadStayAliveMs);
                     if (!taken)
-                        return;
-                    callbackAndState.Callback(callbackAndState.State);
+                    {
+                        lock (_lockObject)
+                        {
+                            if (_tasks.Count == 0)
+                            {
+                                _doingWork = false;
+                                return;
+                            }
+                        }
+
+                        continue;
+                    }
+
+                    try
+                    {
+                        Debug.Assert(_doingWork);
+                        callbackAndState.Callback(callbackAndState.State);
+                    }
+                    catch (Exception)
+                    {
+                        // No logging until 5.0
+                    }
                 }
             }
+            catch (Exception)
+            {
+                // Here we attempt to catch any exception coming from BlockingCollection _tasks
+                lock (_lockObject)
+                    _doingWork = false;
+            }
             finally
             {
-                lock (this) { _thread = null; }
+                Debug.Assert(!_doingWork);
+                _thread = null;
             }
         }
 
