PR

Author: sanych-sun

src/MongoDB.Driver/OperationContext.cs

@@ -17,7 +17,6 @@
 using System.Diagnostics;
 using System.Threading;
 using System.Threading.Tasks;
-using MongoDB.Driver.Core.Misc;
 
 namespace MongoDB.Driver
 {
@@ -26,26 +25,26 @@ internal sealed class OperationContext
         // TODO: this static field is temporary here and will be removed in a future PRs in scope of CSOT.
         public static readonly OperationContext NoTimeout = new(System.Threading.Timeout.InfiniteTimeSpan, CancellationToken.None);
 
-        private readonly Stopwatch _stopwatch;
-
         public OperationContext(TimeSpan timeout, CancellationToken cancellationToken)
             : this(Stopwatch.StartNew(), timeout, cancellationToken)
         {
         }
 
         internal OperationContext(Stopwatch stopwatch, TimeSpan timeout, CancellationToken cancellationToken)
         {
-            _stopwatch = stopwatch;
+            Stopwatch = stopwatch;
             Timeout = timeout;
             CancellationToken = cancellationToken;
         }
 
         public CancellationToken CancellationToken { get; }
 
-        public TimeSpan Elapsed => _stopwatch.Elapsed;
+        public TimeSpan Elapsed => Stopwatch.Elapsed;
 
         public TimeSpan Timeout { get; }
 
+        public OperationContext ParentContext { get; private init; }
+
         public TimeSpan RemainingTimeout
         {
             get
@@ -55,10 +54,12 @@ public TimeSpan RemainingTimeout
                     return System.Threading.Timeout.InfiniteTimeSpan;
                 }
 
-                return Timeout - _stopwatch.Elapsed;
+                return Timeout - Elapsed;
             }
         }
 
+        private Stopwatch Stopwatch { get; }
+
         public bool IsTimedOut()
         {
             var remainingTimeout = RemainingTimeout;
@@ -82,7 +83,10 @@ public OperationContext WithTimeout(TimeSpan timeout)
                 timeout = remainingTimeout;
             }
 
-            return new OperationContext(timeout, CancellationToken);
+            return new OperationContext(timeout, CancellationToken)
+            {
+                ParentContext = this
+            };
         }
 
         public void WaitTask(Task task)

tests/MongoDB.Driver.Tests/OperationContextTests.cs

@@ -35,6 +35,7 @@ public void Constructor_should_initialize_properties()
 
             operationContext.Timeout.Should().Be(timeout);
             operationContext.CancellationToken.Should().Be(cancellationToken);
+            operationContext.ParentContext.Should().BeNull();
         }
 
         [Fact]
@@ -127,6 +128,15 @@ public void WithTimeout_should_calculate_proper_timeout(TimeSpan expected, TimeS
             [TimeSpan.FromMilliseconds(5), TimeSpan.FromMilliseconds(10), TimeSpan.FromMilliseconds(5)],
         ];
 
+        [Fact]
+        public void WithTimeout_should_set_ParentContext()
+        {
+            var operationContext = new OperationContext(new Stopwatch(), Timeout.InfiniteTimeSpan, CancellationToken.None);
+            var resultContext = operationContext.WithTimeout(TimeSpan.FromSeconds(10));
+
+            resultContext.ParentContext.Should().Be(operationContext);
+        }
+
         // TODO: Add tests for WaitTask and WaitTaskAsync.
     }
 }