chore: apply suggestions

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

Author: killme2008

src/GreptimeDB.Ingester/Client/BulkWriter.cs

@@ -28,6 +28,7 @@ public sealed partial class BulkWriter : IBulkWriter
     private Task? _recvTask;
     private uint _serverAffectedRows;
     private volatile Exception? _recvError;
+    private readonly CancellationTokenSource _cts = new();
     private int _completed;
     private int _disposed;
 
@@ -54,6 +55,11 @@ public async ValueTask WriteAsync(Table.Table table, CancellationToken cancellat
 
         if (_recvError != null)
         {
+            if (_recvError is GreptimeException)
+            {
+                throw _recvError;
+            }
+
             throw new GreptimeException($"Stream already failed: {_recvError.Message}", _recvError);
         }
 
@@ -102,11 +108,16 @@ public async ValueTask<uint> CompleteAsync(CancellationToken cancellationToken =
 
             if (_recvTask != null)
             {
-                await _recvTask.ConfigureAwait(false);
+                await _recvTask.WaitAsync(cancellationToken).ConfigureAwait(false);
             }
 
             if (_recvError != null)
             {
+                if (_recvError is GreptimeException)
+                {
+                    throw _recvError;
+                }
+
                 throw new GreptimeException($"Bulk write failed: {_recvError.Message}", _recvError);
             }
 
@@ -132,6 +143,12 @@ public async ValueTask DisposeAsync()
             return;
         }
 
+#if NET8_0_OR_GREATER
+        await _cts.CancelAsync().ConfigureAwait(false);
+#else
+        _cts.Cancel();
+#endif
+
         _putCall?.Dispose();
         _recordBatchBuilder.Dispose();
 
@@ -141,12 +158,14 @@ public async ValueTask DisposeAsync()
             {
                 await _recvTask.WaitAsync(TimeSpan.FromSeconds(5)).ConfigureAwait(false);
             }
-            catch
+            catch (Exception ex)
             {
-                // Ignore exceptions during disposal
+                LogBulkWriteError(_logger, ex.Message);
             }
         }
 
+        _cts.Dispose();
+
         LogBulkWriterDisposed(_logger);
     }
 
@@ -169,16 +188,28 @@ private async Task InitializeStreamAsync(
         _putCall = await _flightClient.StartPut(descriptor, schema, headers, deadline: null, cancellationToken)
             .ConfigureAwait(false);
 
-        _recvTask = DrainResponsesAsync(_putCall.ResponseStream);
+        _recvTask = RunRecvLoopAsync(_putCall.ResponseStream);
 
         LogStreamInitialized(_logger, tableName);
     }
 
-    internal async Task DrainResponsesAsync(IAsyncStreamReader<FlightPutResult> responseStream)
+    private async Task RunRecvLoopAsync(IAsyncStreamReader<FlightPutResult> responseStream)
+    {
+        var (affectedRows, error) = await DrainResponsesAsync(responseStream, _cts.Token).ConfigureAwait(false);
+        _serverAffectedRows = affectedRows;
+        _recvError = error;
+    }
+
+    internal static async Task<(uint AffectedRows, Exception? Error)> DrainResponsesAsync(
+        IAsyncStreamReader<FlightPutResult> responseStream,
+        CancellationToken cancellationToken = default)
     {
+        uint affectedRows = 0;
+        Exception? error = null;
+
         try
         {
-            while (await responseStream.MoveNext(CancellationToken.None).ConfigureAwait(false))
+            while (await responseStream.MoveNext(cancellationToken).ConfigureAwait(false))
             {
                 var result = responseStream.Current;
                 if (result.ApplicationMetadata == null || result.ApplicationMetadata.IsEmpty)
@@ -191,24 +222,30 @@ internal async Task DrainResponsesAsync(IAsyncStreamReader<FlightPutResult> resp
                     var resp = JsonSerializer.Deserialize<DoPutResponse>(result.ApplicationMetadata.Span);
                     if (resp != null)
                     {
-                        _serverAffectedRows += resp.AffectedRows;
+                        affectedRows += resp.AffectedRows;
                     }
                 }
                 catch (JsonException ex)
                 {
-                    _recvError ??= new GreptimeException(
+                    error ??= new GreptimeException(
                         $"Failed to deserialize PutResult metadata: {ex.Message}", ex);
                 }
             }
         }
+        catch (OperationCanceledException ex)
+        {
+            error ??= ex;
+        }
         catch (RpcException ex)
         {
-            _recvError ??= ex;
+            error ??= ex;
         }
-        catch (ObjectDisposedException)
+        catch (ObjectDisposedException ex)
         {
-            // Stream disposed during shutdown, expected
+            error ??= ex;
         }
+
+        return (affectedRows, error);
     }
 
     private void ThrowIfDisposed()

src/GreptimeDB.Ingester/Client/StreamIngestWriter.cs

@@ -150,9 +150,9 @@ public async ValueTask DisposeAsync()
         {
             await _sendTask.WaitAsync(TimeSpan.FromSeconds(5)).ConfigureAwait(false);
         }
-        catch
+        catch (Exception ex)
         {
-            // Ignore exceptions during disposal
+            LogStreamError(_logger, ex.Message);
         }
 
         _cts.Dispose();

tests/GreptimeDB.Ingester.Tests/BulkWriterTests.cs

@@ -1,7 +1,7 @@
 using Apache.Arrow.Flight;
 using FluentAssertions;
-using Google.Protobuf;
 using GreptimeDB.Ingester.Client;
+using GreptimeDB.Ingester.Exceptions;
 using Grpc.Core;
 using Moq;
 using Xunit;
@@ -21,12 +21,11 @@ public async Task DrainResponsesAsync_AccumulatesAffectedRows()
         };
 
         var mockStream = CreateMockStream(responses);
-        var writer = CreateWriter();
 
-        await writer.DrainResponsesAsync(mockStream.Object);
+        var (affectedRows, error) = await BulkWriter.DrainResponsesAsync(mockStream.Object);
 
-        var rows = GetServerAffectedRows(writer);
-        rows.Should().Be(40);
+        affectedRows.Should().Be(40);
+        error.Should().BeNull();
     }
 
     [Fact]
@@ -36,19 +35,12 @@ public async Task DrainResponsesAsync_CapturesRpcException()
         mockStream.Setup(s => s.MoveNext(It.IsAny<CancellationToken>()))
             .ThrowsAsync(new RpcException(new Status(StatusCode.Internal, "server error")));
 
-        var writer = CreateWriter();
+        var (affectedRows, error) = await BulkWriter.DrainResponsesAsync(mockStream.Object);
 
-        await writer.DrainResponsesAsync(mockStream.Object);
-
-        var error = GetRecvError(writer);
+        affectedRows.Should().Be(0);
         error.Should().BeOfType<RpcException>();
     }
 
-    private static BulkWriter CreateWriter()
-    {
-        return new BulkWriter(null!, "test-db", null);
-    }
-
     private static Mock<IAsyncStreamReader<FlightPutResult>> CreateMockStream(FlightPutResult[] responses)
     {
         var mock = new Mock<IAsyncStreamReader<FlightPutResult>>();
@@ -66,18 +58,4 @@ private static Mock<IAsyncStreamReader<FlightPutResult>> CreateMockStream(Flight
 
         return mock;
     }
-
-    private static uint GetServerAffectedRows(BulkWriter writer)
-    {
-        var field = typeof(BulkWriter).GetField("_serverAffectedRows",
-            System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance);
-        return (uint)field!.GetValue(writer)!;
-    }
-
-    private static Exception? GetRecvError(BulkWriter writer)
-    {
-        var field = typeof(BulkWriter).GetField("_recvError",
-            System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance);
-        return (Exception?)field!.GetValue(writer);
-    }
 }