|
1 | 1 | using System.ComponentModel.Composition; |
| 2 | +using Azure; |
2 | 3 | using Azure.Data.Tables; |
3 | 4 | using Azure.Identity; |
4 | 5 | using Cosmos.DataTransfer.AzureTableAPIExtension.Data; |
5 | 6 | using Cosmos.DataTransfer.AzureTableAPIExtension.Settings; |
6 | 7 | using Cosmos.DataTransfer.Interfaces; |
7 | 8 | using Microsoft.Extensions.Configuration; |
8 | 9 | using Microsoft.Extensions.Logging; |
| 10 | +using Polly; |
9 | 11 |
|
10 | 12 | namespace Cosmos.DataTransfer.AzureTableAPIExtension |
11 | 13 | { |
12 | 14 | [Export(typeof(IDataSinkExtension))] |
13 | 15 | public class AzureTableAPIDataSinkExtension : IDataSinkExtensionWithSettings |
14 | 16 | { |
| 17 | + private static readonly int[] TransientStatusCodes = { 408, 429, 500, 502, 503, 504 }; |
| 18 | + |
15 | 19 | public string DisplayName => "AzureTableAPI"; |
16 | 20 |
|
17 | 21 | public async Task WriteAsync(IAsyncEnumerable<IDataItem> dataItems, IConfiguration config, IDataSourceExtension dataSource, ILogger logger, CancellationToken cancellationToken = default) |
@@ -41,21 +45,54 @@ public async Task WriteAsync(IAsyncEnumerable<IDataItem> dataItems, IConfigurati |
41 | 45 |
|
42 | 46 | var tableClient = serviceClient.GetTableClient(settings.Table); |
43 | 47 |
|
44 | | - await tableClient.CreateIfNotExistsAsync(); |
| 48 | + await tableClient.CreateIfNotExistsAsync().ConfigureAwait(false); |
45 | 49 |
|
46 | | - var createTasks = new List<Task>(); |
47 | | - await foreach(var item in dataItems.WithCancellation(cancellationToken)) |
48 | | - { |
49 | | - var entity = item.ToTableEntity(settings.PartitionKeyFieldName, settings.RowKeyFieldName); |
50 | | - createTasks.Add(tableClient.AddEntityAsync(entity)); |
51 | | - } |
| 50 | + var maxConcurrency = settings.MaxConcurrentEntityWrites ?? 10; |
52 | 51 |
|
53 | | - await Task.WhenAll(createTasks); |
| 52 | + logger.LogInformation("Writing data to Azure Table Storage with a maximum of {MaxConcurrency} concurrent writes.", maxConcurrency); |
| 53 | + |
| 54 | + logger.LogInformation("Using PartitionKeyFieldName: `{ParitionKeyFieldName}` and RowKeyFieldName: `{RowKeyFieldName}`", settings.PartitionKeyFieldName, settings.RowKeyFieldName); |
| 55 | + |
| 56 | + await Parallel.ForEachAsync<IDataItem>(dataItems, |
| 57 | + new ParallelOptions { MaxDegreeOfParallelism = maxConcurrency, CancellationToken = cancellationToken }, |
| 58 | + async (item, ct) => |
| 59 | + { |
| 60 | + try |
| 61 | + { |
| 62 | + var entity = item.ToTableEntity(settings.PartitionKeyFieldName, settings.RowKeyFieldName); |
| 63 | + await AddEntityWithRetryAsync(tableClient, entity, ct); |
| 64 | + } |
| 65 | + catch (Exception ex) |
| 66 | + { |
| 67 | + logger.LogError(ex, "Error adding entity to table."); |
| 68 | + } |
| 69 | + }); |
| 70 | + logger.LogInformation("Finished writing data to Azure Table Storage."); |
54 | 71 | } |
55 | 72 |
|
56 | 73 | public IEnumerable<IDataExtensionSettings> GetSettings() |
57 | 74 | { |
58 | 75 | yield return new AzureTableAPIDataSinkSettings(); |
59 | 76 | } |
| 77 | + |
| 78 | + /// <summary> |
| 79 | + /// Adds an entity to the Azure Table Storage with retry logic for transient errors. |
| 80 | + /// This method uses the Polly library to implement a retry policy with exponential backoff. |
| 81 | + /// </summary> |
| 82 | + /// <param name="tableClient"></param> |
| 83 | + /// <param name="entity"></param> |
| 84 | + /// <param name="cancellationToken"></param> |
| 85 | + /// <returns></returns> |
| 86 | + private static async Task AddEntityWithRetryAsync(TableClient tableClient, TableEntity entity, CancellationToken cancellationToken) |
| 87 | + { |
| 88 | + var retryPolicy = Policy |
| 89 | + .Handle<RequestFailedException>(ex => TransientStatusCodes.Contains(ex.Status)) |
| 90 | + .WaitAndRetryAsync(3, retryAttempt => TimeSpan.FromSeconds(Math.Pow(2, retryAttempt))); |
| 91 | + |
| 92 | + await retryPolicy.ExecuteAsync(async () => |
| 93 | + { |
| 94 | + await tableClient.AddEntityAsync(entity, cancellationToken); |
| 95 | + }); |
| 96 | + } |
60 | 97 | } |
61 | 98 | } |
0 commit comments