Fix #3954 make sure BulkAll() is aborted properly if the whole bulk request keeps returning a bad statuscode (#4014)

Ensure BulkAll() is aborted properly if the whole bulk request keeps returning a bad statuscode.

Author: Mpdreamz

src/Nest/CommonAbstractions/Extensions/Extensions.cs

@@ -4,10 +4,12 @@
 using System.Collections.ObjectModel;
 using System.Linq;
 using System.Reflection;
+using System.Runtime.ExceptionServices;
 using System.Runtime.Serialization;
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
+using Elasticsearch.Net;
 using Elasticsearch.Net.Utf8Json.Internal;
 
 namespace Nest
@@ -218,6 +220,12 @@ internal static async Task ForEachAsync<TSource, TResult>(
 						continue;
 
 					var task = await Task.WhenAny(tasks).ConfigureAwait(false);
+					if (task.Exception != null
+						&& (task.IsFaulted && task.Exception.Flatten().InnerExceptions.First() is Exception e))
+					{
+						ExceptionDispatchInfo.Capture(e).Throw();
+						return;
+					}
 					tasks.Remove(task);
 				}
 

src/Nest/Document/Multiple/BulkAll/BulkAllObservable.cs

@@ -22,6 +22,7 @@ public class BulkAllObservable<T> : IDisposable, IObservable<BulkAllResponse> wh
 		private readonly Func<BulkResponseItemBase, T, bool> _retryPredicate;
 		private Action _incrementFailed = () => { };
 		private Action _incrementRetries = () => { };
+		private Action<BulkResponse> _bulkResponseCallback;
 
 		public BulkAllObservable(
 			IElasticClient client,
@@ -36,6 +37,8 @@ public BulkAllObservable(
 			_bulkSize = _partitionedBulkRequest.Size ?? CoordinatedRequestDefaults.BulkAllSizeDefault;
 			_retryPredicate = _partitionedBulkRequest.RetryDocumentPredicate ?? RetryBulkActionPredicate;
 			_droppedDocumentCallBack = _partitionedBulkRequest.DroppedDocumentCallback ?? DroppedDocumentCallbackDefault;
+			_bulkResponseCallback = _partitionedBulkRequest.BulkResponseCallback;
+			
 			_maxDegreeOfParallelism =
 				_partitionedBulkRequest.MaxDegreeOfParallelism ?? CoordinatedRequestDefaults.BulkAllMaxDegreeOfParallelismDefault;
 			_compositeCancelTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
@@ -125,6 +128,8 @@ private async Task<BulkAllResponse> BulkAsync(IList<T> buffer, long page, int ba
 				.ConfigureAwait(false);
 
 			_compositeCancelToken.ThrowIfCancellationRequested();
+			
+			_bulkResponseCallback?.Invoke(response);
 
 			if (!response.ApiCall.Success)
 				return await HandleBulkRequest(buffer, page, backOffRetries, response).ConfigureAwait(false);
@@ -167,28 +172,39 @@ private void HandleDroppedDocuments(List<Tuple<BulkResponseItemBase, T>> dropped
 		private async Task<BulkAllResponse> HandleBulkRequest(IList<T> buffer, long page, int backOffRetries, BulkResponse response)
 		{
 			var clientException = response.ApiCall.OriginalException as ElasticsearchClientException;
-			var failureReason = clientException?.FailureReason.GetValueOrDefault(PipelineFailure.Unexpected);
+			var failureReason = clientException?.FailureReason; 
+			var reason = failureReason?.GetStringValue() ?? nameof(PipelineFailure.BadRequest);
 			switch (failureReason)
 			{
 				case PipelineFailure.MaxRetriesReached:
 					if (response.ApiCall.AuditTrail.Last().Event == AuditEvent.FailedOverAllNodes)
 						throw ThrowOnBadBulk(response, $"BulkAll halted after attempted bulk failed over all the active nodes");
 
+					ThrowOnExhaustedRetries();
 					return await RetryDocuments(page, ++backOffRetries, buffer).ConfigureAwait(false);
 				case PipelineFailure.CouldNotStartSniffOnStartup:
 				case PipelineFailure.BadAuthentication:
 				case PipelineFailure.NoNodesAttempted:
 				case PipelineFailure.SniffFailure:
 				case PipelineFailure.Unexpected:
-					throw ThrowOnBadBulk(response,
-						$"BulkAll halted after {nameof(PipelineFailure)}{failureReason.GetStringValue()} from _bulk");
+					throw ThrowOnBadBulk(response, $"BulkAll halted after {nameof(PipelineFailure)}.{reason} from _bulk");
 				case PipelineFailure.BadResponse:
 				case PipelineFailure.PingFailure:
 				case PipelineFailure.MaxTimeoutReached:
 				case PipelineFailure.BadRequest:
 				default:
+					ThrowOnExhaustedRetries();
 					return await RetryDocuments(page, ++backOffRetries, buffer).ConfigureAwait(false);
 			}
+
+			void ThrowOnExhaustedRetries()
+			{
+				if (_partitionedBulkRequest.ContinueAfterDroppedDocuments || backOffRetries < _backOffRetries) return;
+
+				throw ThrowOnBadBulk(response,
+					$"BulkAll halted after {nameof(PipelineFailure)}.{reason} from _bulk and exhausting retries ({backOffRetries})"
+				);
+			}
 		}
 
 		private async Task<BulkAllResponse> RetryDocuments(long page, int backOffRetries, IList<T> retryDocuments)

src/Nest/Document/Multiple/BulkAll/BulkAllRequest.cs

@@ -85,6 +85,12 @@ public interface IBulkAllRequest<T> where T : class
 		/// non-negative value less than or equal to the total number of copies for the shard (number of replicas + 1)
 		/// </summary>
 		int? WaitForActiveShards { get; set; }
+		
+		/// <summary>
+		/// Be notified every time a bulk response returns, this includes retries.
+		/// <see cref="IObserver{T}.OnNext"/> is only called for successful batches.
+		/// </summary>
+		Action<BulkResponse> BulkResponseCallback { get; set; }
 	}
 
 	public class BulkAllRequest<T> : IBulkAllRequest<T>
@@ -146,6 +152,9 @@ public BulkAllRequest(IEnumerable<T> documents)
 
 		/// <inheritdoc />
 		public int? WaitForActiveShards { get; set; }
+		
+		/// <inheritdoc />
+		public Action<BulkResponse> BulkResponseCallback { get; set; }
 	}
 
 	public class BulkAllDescriptor<T> : DescriptorBase<BulkAllDescriptor<T>, IBulkAllRequest<T>>, IBulkAllRequest<T>
@@ -177,6 +186,7 @@ public BulkAllDescriptor(IEnumerable<T> documents)
 		int? IBulkAllRequest<T>.Size { get; set; }
 		Time IBulkAllRequest<T>.Timeout { get; set; }
 		int? IBulkAllRequest<T>.WaitForActiveShards { get; set; }
+		Action<BulkResponse> IBulkAllRequest<T>.BulkResponseCallback { get; set; }
 
 		/// <inheritdoc cref="IBulkAllRequest{T}.MaxDegreeOfParallelism" />
 		public BulkAllDescriptor<T> MaxDegreeOfParallelism(int? parallelism) =>
@@ -238,5 +248,9 @@ public BulkAllDescriptor<T> BackPressure(int maxConcurrency, int? backPressureFa
 		/// <inheritdoc cref="IBulkAllRequest{T}.DroppedDocumentCallback" />
 		public BulkAllDescriptor<T> DroppedDocumentCallback(Action<BulkResponseItemBase, T> callback) =>
 			Assign(callback, (a, v) => a.DroppedDocumentCallback = v);
+		
+		/// <inheritdoc cref="IBulkAllRequest{T}.BulkResponseCallback" />
+		public BulkAllDescriptor<T> BulkResponseCallback(Action<BulkResponse> callback) =>
+			Assign(callback, (a, v) => a.BulkResponseCallback = v);
 	}
 }

src/Tests/Tests/Document/Multiple/BulkAll/BulkAllExceptionApiTests.cs

@@ -1,9 +1,13 @@
 using System;
 using System.Threading;
 using Elastic.Xunit.XunitPlumbing;
+using Elasticsearch.Net;
 using FluentAssertions;
 using Nest;
+using Tests.Core.Client.Settings;
 using Tests.Core.ManagedElasticsearch.Clusters;
+using Tests.Domain.Extensions;
+using Tests.Framework.VirtualClustering;
 
 namespace Tests.Document.Multiple.BulkAll
 {
@@ -48,4 +52,65 @@ [I] public void WaitBulkAllThrowsAndIsCaught()
 			}
 		}
 	}
+
+
+	public class BulkAllBadRetriesApiTests : BulkAllApiTestsBase
+	{
+		public BulkAllBadRetriesApiTests(IntrusiveOperationCluster cluster) : base(cluster) { }
+		
+		[U] public void Completes()
+		{
+			var client = VirtualClusterWith.Nodes(2)
+				.ClientCalls(c => c.FailAlways())
+				.StaticConnectionPool()
+				.AllDefaults()
+				.Client;
+			
+			
+			var index = CreateIndexName();
+
+			var size = 1000;
+			var pages = 10;
+			var seenPages = 0;
+			var numberOfDocuments = size * pages;
+			var documents = CreateLazyStreamOfDocuments(numberOfDocuments);
+			var requests = 0;
+
+			Exception ex = null;
+			var tokenSource = new CancellationTokenSource();
+			var observableBulk = client.BulkAll(documents, f => f
+					.MaxDegreeOfParallelism(1)
+					.BulkResponseCallback(r => Interlocked.Increment(ref requests))
+					.BackOffTime(TimeSpan.FromMilliseconds(1))
+					.BackOffRetries(2)
+					.Size(size)
+					.RefreshOnCompleted()
+					.Index(index)
+					.BufferToBulk((r, buffer) => r.IndexMany(buffer))
+				, tokenSource.Token);
+
+			try
+			{
+				observableBulk.Wait(TimeSpan.FromSeconds(30), b =>
+				{
+					Interlocked.Increment(ref seenPages);
+				});
+			}
+			catch (Exception e)
+			{
+				ex = e;
+			}
+			ex.Should().NotBeNull();
+
+			var clientException = ex.Should().BeOfType<ElasticsearchClientException>().Subject;
+			
+			clientException.Message.Should()
+				.StartWith("BulkAll halted after")
+				.And.EndWith("from _bulk and exhausting retries (2)");
+			
+			requests.Should().Be(3);
+			// OnNext only called for successful batches.
+			seenPages.Should().Be(0);
+		}
+	}
 }

src/Tests/Tests/Framework/VirtualClustering/VirtualizedCluster.cs

@@ -30,7 +30,7 @@ public VirtualizedCluster(TestableDateTimeProvider dateTimeProvider, ConnectionS
 		}
 
 		public IConnectionPool ConnectionPool => Client.ConnectionSettings.ConnectionPool;
-		private ElasticClient Client => _fixedRequestPipeline?.Client;
+		public ElasticClient Client => _fixedRequestPipeline?.Client;
 
 		public VirtualizedCluster ClientProxiesTo(
 			Func<IElasticClient, Func<RequestConfigurationDescriptor, IRequestConfiguration>, IResponse> sync,