Better handling of Future handles

FoundationDB.Client/Core/IFdbTransactionHandler.cs

@@ -52,6 +52,7 @@ public interface IFdbTransactionHandler : IDisposable
 		void SetOption(FdbTransactionOption option, Slice data);
 
 		/// <summary>Returns this transaction snapshot read version.</summary>
+		/// <param name="ct">Token used to cancel the operation from the outside, if different than the cancellation token of the transaction itself</param>
 		Task<long> GetReadVersionAsync(CancellationToken ct);
 
 		/// <summary>Retrieves the database version number at which a given transaction was committed.</summary>

FoundationDB.Client/FdbTransaction.Snapshot.cs

@@ -101,7 +101,7 @@ public Task<Slice> GetAsync(Slice key)
 				if (Logging.On && Logging.IsVerbose) Logging.Verbose(this, "GetAsync", $"Getting value for '{key.ToString()}'");
 #endif
 
-				return m_parent.m_handler.GetAsync(key, snapshot: true, ct: m_parent.m_cancellation);
+				return m_parent.m_handler.GetAsync(key, snapshot: true, ct: CancellationToken.None);
 			}
 
 			public Task<Slice[]> GetValuesAsync(Slice[] keys)
@@ -116,7 +116,7 @@ public Task<Slice[]> GetValuesAsync(Slice[] keys)
 				if (Logging.On && Logging.IsVerbose) Logging.Verbose(this, "GetValuesAsync", $"Getting batch of {keys.Length} values ...");
 #endif
 
-				return m_parent.m_handler.GetValuesAsync(keys, snapshot: true, ct: m_parent.m_cancellation);
+				return m_parent.m_handler.GetValuesAsync(keys, snapshot: true, ct: CancellationToken.None);
 			}
 
 			public async Task<Slice> GetKeyAsync(KeySelector selector)
@@ -129,7 +129,7 @@ public async Task<Slice> GetKeyAsync(KeySelector selector)
 				if (Logging.On && Logging.IsVerbose) Logging.Verbose(this, "GetKeyAsync", $"Getting key '{selector.ToString()}'");
 #endif
 
-				var key = await m_parent.m_handler.GetKeyAsync(selector, snapshot: true, ct: m_parent.m_cancellation).ConfigureAwait(false);
+				var key = await m_parent.m_handler.GetKeyAsync(selector, snapshot: true, ct: CancellationToken.None).ConfigureAwait(false);
 
 				// don't forget to truncate keys that would fall outside of the database's globalspace !
 				return m_parent.m_database.BoundCheck(key);
@@ -149,7 +149,7 @@ public Task<Slice[]> GetKeysAsync(KeySelector[] selectors)
 				if (Logging.On && Logging.IsVerbose) Logging.Verbose(this, "GetKeysCoreAsync", $"Getting batch of {selectors.Length} keys ...");
 #endif
 
-				return m_parent.m_handler.GetKeysAsync(selectors, snapshot: true, ct: m_parent.m_cancellation);
+				return m_parent.m_handler.GetKeysAsync(selectors, snapshot: true, ct: CancellationToken.None);
 			}
 
 			public Task<FdbRangeChunk> GetRangeAsync(KeySelector beginInclusive, KeySelector endExclusive, FdbRangeOptions options, int iteration)
@@ -165,7 +165,7 @@ public Task<FdbRangeChunk> GetRangeAsync(KeySelector beginInclusive, KeySelector
 				// The iteration value is only needed when in iterator mode, but then it should start from 1
 				if (iteration == 0) iteration = 1;
 
-				return m_parent.m_handler.GetRangeAsync(beginInclusive, endExclusive, options, iteration, snapshot: true, ct: m_parent.m_cancellation);
+				return m_parent.m_handler.GetRangeAsync(beginInclusive, endExclusive, options, iteration, snapshot: true, ct: CancellationToken.None);
 			}
 
 			public FdbRangeQuery<KeyValuePair<Slice, Slice>> GetRange(KeySelector beginInclusive, KeySelector endExclusive, FdbRangeOptions options)
@@ -178,7 +178,7 @@ public FdbRangeQuery<KeyValuePair<Slice, Slice>> GetRange(KeySelector beginInclu
 			public Task<string[]> GetAddressesForKeyAsync(Slice key)
 			{
 				EnsureCanRead();
-				return m_parent.m_handler.GetAddressesForKeyAsync(key, ct: m_parent.m_cancellation);
+				return m_parent.m_handler.GetAddressesForKeyAsync(key, CancellationToken.None);
 			}
 
 			void IFdbReadOnlyTransaction.Cancel()

FoundationDB.Client/FdbTransaction.cs

@@ -93,6 +93,9 @@ public sealed partial class FdbTransaction : IFdbTransaction
 		/// <summary>Random token (but constant per transaction retry) used to generate incomplete VersionStamps</summary>
 		private ulong m_versionStampToken;
 
+		/// <summary>Used to cancel the transaction if the parent CTS fires</summary>
+		private CancellationTokenRegistration m_ctr;
+
 		#endregion
 
 		#region Constructors...
@@ -111,6 +114,15 @@ internal FdbTransaction(FdbDatabase db, FdbOperationContext context, int id, IFd
 
 			m_readOnly = (mode & FdbTransactionMode.ReadOnly) != 0;
 			m_handler = handler;
+
+			if (m_cancellation.IsCancellationRequested)
+			{ // already dead?
+				Cancel(explicitly: false);
+			}
+			else
+			{
+				m_ctr = m_cancellation.RegisterWithoutEC(CancellationHandler, this);
+			}
 		}
 
 		#endregion
@@ -263,7 +275,7 @@ public Task<long> GetReadVersionAsync()
 			// can be called after the transaction has been committed
 			EnsureCanRetry();
 
-			return m_handler.GetReadVersionAsync(m_cancellation);
+			return m_handler.GetReadVersionAsync(CancellationToken.None);
 		}
 
 		/// <summary>Retrieves the database version number at which a given transaction was committed.</summary>
@@ -386,7 +398,7 @@ public Task<Slice> GetAsync(Slice key)
 			if (Logging.On && Logging.IsVerbose) Logging.Verbose(this, "GetAsync", $"Getting value for '{key.ToString()}'");
 #endif
 
-			return m_handler.GetAsync(key, snapshot: false, ct: m_cancellation);
+			return m_handler.GetAsync(key, snapshot: false, ct: CancellationToken.None);
 		}
 
 		#endregion
@@ -409,7 +421,7 @@ public Task<Slice[]> GetValuesAsync(Slice[] keys)
 			if (Logging.On && Logging.IsVerbose) Logging.Verbose(this, "GetValuesAsync", $"Getting batch of {keys.Length} values ...");
 #endif
 
-			return m_handler.GetValuesAsync(keys, snapshot: false, ct: m_cancellation);
+			return m_handler.GetValuesAsync(keys, snapshot: false, ct: CancellationToken.None);
 		}
 
 		#endregion
@@ -439,7 +451,7 @@ public Task<FdbRangeChunk> GetRangeAsync(KeySelector beginInclusive, KeySelector
 			// The iteration value is only needed when in iterator mode, but then it should start from 1
 			if (iteration == 0) iteration = 1;
 
-			return m_handler.GetRangeAsync(beginInclusive, endExclusive, options, iteration, snapshot: false, ct: m_cancellation);
+			return m_handler.GetRangeAsync(beginInclusive, endExclusive, options, iteration, snapshot: false, ct: CancellationToken.None);
 		}
 
 		#endregion
@@ -492,7 +504,7 @@ public async Task<Slice> GetKeyAsync(KeySelector selector)
 			if (Logging.On && Logging.IsVerbose) Logging.Verbose(this, "GetKeyAsync", $"Getting key '{selector.ToString()}'");
 #endif
 
-			var key = await m_handler.GetKeyAsync(selector, snapshot: false, ct: m_cancellation).ConfigureAwait(false);
+			var key = await m_handler.GetKeyAsync(selector, snapshot: false, ct: CancellationToken.None).ConfigureAwait(false);
 
 			// don't forget to truncate keys that would fall outside of the database's globalspace !
 			return m_database.BoundCheck(key);
@@ -520,7 +532,7 @@ public Task<Slice[]> GetKeysAsync(KeySelector[] selectors)
 			if (Logging.On && Logging.IsVerbose) Logging.Verbose(this, "GetKeysAsync", $"Getting batch of {selectors.Length} keys ...");
 #endif
 
-			return m_handler.GetKeysAsync(selectors, snapshot: false, ct: m_cancellation);
+			return m_handler.GetKeysAsync(selectors, snapshot: false, ct: CancellationToken.None);
 		}
 
 		#endregion
@@ -727,7 +739,7 @@ public Task<string[]> GetAddressesForKeyAsync(Slice key)
 			if (Logging.On && Logging.IsVerbose) Logging.Verbose(this, "GetAddressesForKeyAsync", $"Getting addresses for key '{FdbKey.Dump(key)}'");
 #endif
 
-			return m_handler.GetAddressesForKeyAsync(key, ct: m_cancellation);
+			return m_handler.GetAddressesForKeyAsync(key, CancellationToken.None);
 		}
 
 		#endregion
@@ -750,7 +762,7 @@ public async Task CommitAsync()
 			//TODO: need a STATE_COMMITTING ?
 			try
 			{
-				await m_handler.CommitAsync(m_cancellation).ConfigureAwait(false);
+				await m_handler.CommitAsync(CancellationToken.None).ConfigureAwait(false);
 
 				if (Interlocked.CompareExchange(ref m_state, STATE_COMMITTED, STATE_READY) == STATE_READY)
 				{
@@ -817,7 +829,7 @@ public async Task OnErrorAsync(FdbError code)
 		{
 			EnsureCanRetry();
 
-			await m_handler.OnErrorAsync(code, ct: m_cancellation).ConfigureAwait(false);
+			await m_handler.OnErrorAsync(code, CancellationToken.None).ConfigureAwait(false);
 
 			// If fdb_transaction_on_error succeeds, that means that the transaction has been reset and is usable again
 			var state = this.State;
@@ -875,19 +887,47 @@ public void Reset()
 
 		/// <summary>Rollback this transaction, and dispose it. It should not be used after that.</summary>
 		public void Cancel()
+		{
+			Cancel(explicitly: true);
+		}
+
+		private void Cancel(bool explicitly)
 		{
 			var state = Interlocked.CompareExchange(ref m_state, STATE_CANCELED, STATE_READY);
 			if (state != STATE_READY)
 			{
-				switch(state)
+				if (explicitly)
 				{
-					case STATE_CANCELED: return; // already the case !
+					switch (state)
+					{
+						case STATE_CANCELED:
+						{
+							return; // already the case!
+						}
+						case STATE_COMMITTED:
+						{
+							throw new InvalidOperationException("Cannot cancel transaction that has already been committed");
+						}
+						case STATE_FAILED:
+						{
+							throw new InvalidOperationException("Cannot cancel transaction because it is in a failed state");
+						}
+						case STATE_DISPOSED:
+						{
+							throw new ObjectDisposedException("FdbTransaction", "Cannot cancel transaction because it already has been disposed");
+						}
+						default:
+						{
+							throw new InvalidOperationException(String.Format("Cannot cancel transaction because it is in unknown state {0}", state));
+						}
+					}
+				}
 
-					case STATE_COMMITTED: throw new InvalidOperationException("Cannot cancel transaction that has already been committed");
-					case STATE_FAILED: throw new InvalidOperationException("Cannot cancel transaction because it is in a failed state");
-					case STATE_DISPOSED: throw new ObjectDisposedException("FdbTransaction", "Cannot cancel transaction because it already has been disposed");
-					default: throw new InvalidOperationException($"Cannot cancel transaction because it is in unknown state {state}");
+				if (state == STATE_CANCELED || state == STATE_DISPOSED)
+				{ // it's too late
+					return;
 				}
+
 			}
 
 			if (Logging.On && Logging.IsVerbose) Logging.Verbose(this, "Cancel", "Canceling transaction...");
@@ -897,6 +937,16 @@ public void Cancel()
 			if (Logging.On && Logging.IsVerbose) Logging.Verbose(this, "Cancel", "Transaction has been canceled");
 		}
 
+		private static readonly Action<object> CancellationHandler = CancellationCallback;
+
+		/// <summary>Handler called when the cancellation source of the transaction fires</summary>
+		private static void CancellationCallback(object state)
+		{
+			Contract.Requires(state != null);
+			var trans = (FdbTransaction) state;
+			trans.Cancel(explicitly: false);
+		}
+
 		#endregion
 
 		#region IDisposable...
@@ -1014,6 +1064,7 @@ public void Dispose()
 			{
 				try
 				{
+					m_ctr.Dispose();
 					this.Database.UnregisterTransaction(this);
 					m_cts.SafeCancelAndDispose();
 

FoundationDB.Client/FdbWatch.cs

@@ -82,10 +82,13 @@ public TaskAwaiter<Slice> GetAwaiter()
 
 			if (m_future != null)
 			{
+#if REFACTORING_IN_PROGRESS
 				if (m_future.HasFlag(FdbFuture.Flags.DISPOSED))
 				{
 					throw new ObjectDisposedException("Cannot await a watch that has already been disposed");
 				}
+
+#endif
 				return m_future.Task.GetAwaiter();
 			}
 			throw new InvalidOperationException("Cannot await an empty watch");
@@ -100,7 +103,11 @@ public void Cancel()
 		/// <summary>Dispose the resources allocated by the watch.</summary>
 		public void Dispose()
 		{
-			m_future?.Dispose();
+			if (m_future != null)
+			{
+				//TODO: what should be do? (=> cancel the future?)
+				//m_future.Dispose();
+			}
 		}
 
 		public override string ToString()