From 0fb2e8d299617a3312e1d7a1c9dda8ea1aa6187f Mon Sep 17 00:00:00 2001
From: sakno <roman.sakno@gmail.com>
Date: Tue, 23 Jan 2024 02:29:17 +0200
Subject: [PATCH] Fixed cleanup of async state

---
 .../IO/FileBufferingWriter.Utils.cs           | 50 +++++++++++--------
 src/DotNext.IO/IO/FileWriter.Utils.cs         | 34 +++++++------
 2 files changed, 47 insertions(+), 37 deletions(-)

diff --git a/src/DotNext.IO/IO/FileBufferingWriter.Utils.cs b/src/DotNext.IO/IO/FileBufferingWriter.Utils.cs
index 2ba1fa672..ba24de421 100644
--- a/src/DotNext.IO/IO/FileBufferingWriter.Utils.cs
+++ b/src/DotNext.IO/IO/FileBufferingWriter.Utils.cs
@@ -58,30 +58,38 @@ private void GetAsyncResult(short token)
 
     private void OnWrite()
     {
+        var awaiter = this.awaiter;
+        this.awaiter = default;
+
+        var secondBuffer = this.secondBuffer;
+        this.secondBuffer = default;
+
         try
         {
             awaiter.GetResult();
 
             filePosition += secondBuffer.Length + position;
             position = 0;
-
-            source.SetResult(0);
         }
         catch (Exception e)
         {
             source.SetException(e);
+            return;
         }
-        finally
-        {
-            awaiter = default;
-            secondBuffer = default;
-        }
+
+        source.SetResult(0);
     }
 
     private void OnWriteAndFlush()
     {
         Debug.Assert(fileBackend is not null);
 
+        var awaiter = this.awaiter;
+        this.awaiter = default;
+
+        var secondBuffer = this.secondBuffer;
+        this.secondBuffer = default;
+
         try
         {
             awaiter.GetResult();
@@ -89,22 +97,24 @@ private void OnWriteAndFlush()
             filePosition += secondBuffer.Length + position;
             position = 0;
             RandomAccess.FlushToDisk(fileBackend);
-
-            source.SetResult(0);
         }
         catch (Exception e)
         {
             source.SetException(e);
+            return;
         }
-        finally
-        {
-            awaiter = default;
-            secondBuffer = default;
-        }
+
+        source.SetResult(0);
     }
 
     private void OnWriteAndCopy()
     {
+        var awaiter = this.awaiter;
+        this.awaiter = default;
+
+        var secondBuffer = this.secondBuffer;
+        this.secondBuffer = default;
+
         try
         {
             awaiter.GetResult();
@@ -112,24 +122,20 @@ private void OnWriteAndCopy()
             filePosition += position;
             secondBuffer.CopyTo(buffer.Memory);
             position = secondBuffer.Length;
-
-            source.SetResult(0);
         }
         catch (Exception e)
         {
             source.SetException(e);
+            return;
         }
-        finally
-        {
-            awaiter = default;
-            secondBuffer = default;
-        }
+
+        source.SetResult(0);
     }
 
     private ValueTask Submit(ValueTask task, Action callback)
     {
         awaiter = task.ConfigureAwait(false).GetAwaiter();
-        if (task.IsCompleted)
+        if (awaiter.IsCompleted)
         {
             callback();
         }
diff --git a/src/DotNext.IO/IO/FileWriter.Utils.cs b/src/DotNext.IO/IO/FileWriter.Utils.cs
index 89090ae81..ad8bfbf8a 100644
--- a/src/DotNext.IO/IO/FileWriter.Utils.cs
+++ b/src/DotNext.IO/IO/FileWriter.Utils.cs
@@ -57,28 +57,36 @@ private void GetAsyncResult(short token)
 
     private void OnWrite()
     {
+        var awaiter = this.awaiter;
+        this.awaiter = default;
+
+        var secondBuffer = this.secondBuffer;
+        this.secondBuffer = default;
+
         try
         {
             awaiter.GetResult();
 
             fileOffset += secondBuffer.Length + bufferOffset;
             bufferOffset = 0;
-
-            source.SetResult(0);
         }
         catch (Exception e)
         {
             source.SetException(e);
+            return;
         }
-        finally
-        {
-            awaiter = default;
-            secondBuffer = default;
-        }
+
+        source.SetResult(0);
     }
 
     private void OnWriteAndCopy()
     {
+        var awaiter = this.awaiter;
+        this.awaiter = default;
+
+        var secondBuffer = this.secondBuffer;
+        this.secondBuffer = default;
+
         try
         {
             awaiter.GetResult();
@@ -86,24 +94,20 @@ private void OnWriteAndCopy()
             fileOffset += bufferOffset;
             secondBuffer.CopyTo(buffer.Memory);
             bufferOffset = secondBuffer.Length;
-
-            source.SetResult(0);
         }
         catch (Exception e)
         {
             source.SetException(e);
+            return;
         }
-        finally
-        {
-            awaiter = default;
-            secondBuffer = default;
-        }
+
+        source.SetResult(0);
     }
 
     private ValueTask Submit(ValueTask task, Action callback)
     {
         awaiter = task.ConfigureAwait(false).GetAwaiter();
-        if (task.IsCompleted)
+        if (awaiter.IsCompleted)
         {
             callback();
         }