From fecfa94433293c9874ec312af66e25935631914d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Christian=20K=C3=B6llner?= <christian.koellner.80@gmail.com>
Date: Tue, 20 Aug 2019 21:26:13 +0200
Subject: [PATCH] Troubleshooting Embargo / EmbargoServer TC

---
 Capnp.Net.Runtime.Tests/TcpRpcInterop.cs      |  3 +++
 .../Rpc/RemoteAnswerCapability.cs             | 27 +++++++++++--------
 Capnp.Net.Runtime/Rpc/RpcEngine.cs            | 10 +++----
 Capnp.Net.Runtime/Rpc/Skeleton.cs             | 11 ++------
 4 files changed, 24 insertions(+), 27 deletions(-)

diff --git a/Capnp.Net.Runtime.Tests/TcpRpcInterop.cs b/Capnp.Net.Runtime.Tests/TcpRpcInterop.cs
index cc17f0a..e7403cb 100644
--- a/Capnp.Net.Runtime.Tests/TcpRpcInterop.cs
+++ b/Capnp.Net.Runtime.Tests/TcpRpcInterop.cs
@@ -774,6 +774,7 @@ namespace Capnp.Net.Runtime.Tests
 
                         var cap = new TestCallOrderImpl();
                         cap.CountToDispose = 6;
+                        Skeleton.BeginAssertNotDisposed(cap);
                         var earlyCall = main.GetCallSequence(0, default);
 
                         var echo = main.Echo(cap, default);
@@ -807,6 +808,8 @@ namespace Capnp.Net.Runtime.Tests
                                 Assert.AreEqual(3u, call3.Result);
                                 Assert.AreEqual(4u, call4.Result);
                                 Assert.AreEqual(5u, call5.Result);
+
+                                Skeleton.EndAssertNotDisposed(cap);
                             }
                         }
                     }
diff --git a/Capnp.Net.Runtime/Rpc/RemoteAnswerCapability.cs b/Capnp.Net.Runtime/Rpc/RemoteAnswerCapability.cs
index 61af0e6..816a8a3 100644
--- a/Capnp.Net.Runtime/Rpc/RemoteAnswerCapability.cs
+++ b/Capnp.Net.Runtime/Rpc/RemoteAnswerCapability.cs
@@ -24,6 +24,8 @@ namespace Capnp.Rpc
         {
             _question = question ?? throw new ArgumentNullException(nameof(question));
             _access = access ?? throw new ArgumentNullException(nameof(access));
+
+            _ = AwaitWhenResolved();
         }
 
         async void ReAllowFinishWhenDone(Task task)
@@ -57,21 +59,24 @@ namespace Capnp.Rpc
         {
             get
             {
-                if (_resolvedCap == null && !_question.IsTailCall && _question.IsReturned)
+                lock (_question.ReentrancyBlocker)
                 {
-                    DeserializerState result;
-                    try
+                    if (_resolvedCap == null && !_question.IsTailCall && _question.IsReturned)
                     {
-                        result = _question.WhenReturned.Result;
-                    }
-                    catch (AggregateException exception)
-                    {
-                        throw exception.InnerException;
-                    }
+                        DeserializerState result;
+                        try
+                        {
+                            result = _question.WhenReturned.Result;
+                        }
+                        catch (AggregateException exception)
+                        {
+                            throw exception.InnerException;
+                        }
 
-                    _resolvedCap = new Proxy(_access.Eval(result));
+                        _resolvedCap = new Proxy(_access.Eval(result));
+                    }
+                    return _resolvedCap;
                 }
-                return _resolvedCap;
             }
         }
 
diff --git a/Capnp.Net.Runtime/Rpc/RpcEngine.cs b/Capnp.Net.Runtime/Rpc/RpcEngine.cs
index 2333454..cb20bf6 100644
--- a/Capnp.Net.Runtime/Rpc/RpcEngine.cs
+++ b/Capnp.Net.Runtime/Rpc/RpcEngine.cs
@@ -973,7 +973,7 @@ namespace Capnp.Rpc
 
             void ReleaseExport(uint id, uint count)
             {
-                bool exists, badrc = false;
+                bool exists;
 
                 lock (_reentrancyBlocker)
                 {
@@ -993,9 +993,9 @@ namespace Capnp.Rpc
                                 _revExportTable.ReplacementTryRemove(rc.Cap, out uint _);
                             }
                         }
-                        catch (System.Exception)
+                        catch (System.Exception exception)
                         {
-                            badrc = true;
+                            Logger.LogWarning($"Attempting to release capability with invalid reference count: {exception.Message}");
                         }
                     }
                 }
@@ -1004,10 +1004,6 @@ namespace Capnp.Rpc
                 {
                     Logger.LogWarning("Attempting to release unknown capability ID");
                 }
-                else if (badrc)
-                {
-                    Logger.LogWarning("Attempting to release capability with invalid reference count");
-                }
             }
 
             void ProcessRelease(Release.READER release)
diff --git a/Capnp.Net.Runtime/Rpc/Skeleton.cs b/Capnp.Net.Runtime/Rpc/Skeleton.cs
index 283ce5f..2a098f9 100644
--- a/Capnp.Net.Runtime/Rpc/Skeleton.cs
+++ b/Capnp.Net.Runtime/Rpc/Skeleton.cs
@@ -29,9 +29,7 @@ namespace Capnp.Rpc
             }
         }
 
-#if DEBUG
         const int NoDisposeFlag = 0x4000000;
-#endif
 
         static readonly ConditionalWeakTable<object, Skeleton> _implMap =
             new ConditionalWeakTable<object, Skeleton>();
@@ -63,7 +61,6 @@ namespace Capnp.Rpc
             return new SkeletonRelinquisher(GetOrCreateSkeleton(impl, true));
         }
 
-#if DEBUG
         /// <summary>
         /// This DEBUG-only diagnostic method states that the Skeleton corresponding to a given capability is not expected to
         /// be disposed until the next call to EndAssertNotDisposed().
@@ -84,7 +81,6 @@ namespace Capnp.Rpc
         {
             GetOrCreateSkeleton(impl, false).EndAssertNotDisposed();
         }
-#endif
 
         int _refCount = 0;
 
@@ -103,7 +99,6 @@ namespace Capnp.Rpc
             Interlocked.Increment(ref _refCount);
         }
 
-#if DEBUG
         internal void BeginAssertNotDisposed()
         {
             if ((Interlocked.Add(ref _refCount, NoDisposeFlag) & NoDisposeFlag) == 0)
@@ -118,7 +113,6 @@ namespace Capnp.Rpc
                 throw new InvalidOperationException("Flag already cleared. State is now broken.");
             }
         }
-#endif
 
         internal void Relinquish()
         {
@@ -126,9 +120,8 @@ namespace Capnp.Rpc
 
             if (0 == count)
             {
-#if DEBUG
-                Debug.Assert((_refCount & NoDisposeFlag) == 0, "Skeleton disposal not expected in this state");
-#endif
+                if ((count & NoDisposeFlag) != 0)
+                    throw new InvalidOperationException("Unexpected Skeleton disposal");
 
                 Dispose(true);
                 GC.SuppressFinalize(this);