feat: EFCore supports YdbRetryPolicy (#546)

Author: KirillKurdyukov

examples/EntityFrameworkCore.Ydb.Yandex.Cloud/Program.cs

@@ -12,8 +12,8 @@ await Parser.Default.ParseArguments<CmdOptions>(args).WithParsedAsync(async cmd
 
     var options = new DbContextOptionsBuilder<AppDbContext>()
         .UseYdb(cmd.ConnectionString, builder => builder
-            .WithCredentialsProvider(saProvider)
-            .WithServerCertificates(YcCerts.GetYcServerCertificates())
+            .UseCredentialsProvider(saProvider)
+            .UseServerCertificates(YcCerts.GetYcServerCertificates())
         )
         .Options;
 

slo/src/EF/SloTableContext.cs

@@ -12,7 +12,8 @@ public class SloTableContext : SloTableContext<PooledDbContextFactory<TableDbCon
     protected override string Job => "EF";
 
     protected override PooledDbContextFactory<TableDbContext> CreateClient(Config config) =>
-        new(new DbContextOptionsBuilder<TableDbContext>().UseYdb(config.ConnectionString).Options);
+        new(new DbContextOptionsBuilder<TableDbContext>().UseYdb(config.ConnectionString,
+            builder => builder.EnableRetryIdempotence()).Options);
 
     protected override async Task Create(
         PooledDbContextFactory<TableDbContext> client,

src/EFCore.Ydb/CHANGELOG.md

@@ -1,3 +1,7 @@
+- Added support for the YDB retry policy (ADO.NET) and new configuration methods in `YdbDbContextOptionsBuilder`:
+  - `EnableRetryIdempotence()`: enables retries for errors classified as idempotent. You must ensure the operation itself is idempotent.
+  - `UseRetryPolicy(YdbRetryPolicyConfig retryPolicyConfig)`: configures custom backoff parameters and the maximum number of retry attempts.
+
 ## v0.1.0
 
 - Fixed bug: incompatible coalesce types ([#531](https://github.com/ydb-platform/ydb-dotnet-sdk/issues/531)).

src/EFCore.Ydb/src/Infrastructure/Internal/YdbOptionsExtension.cs

@@ -13,8 +13,6 @@ public class YdbOptionsExtension : RelationalOptionsExtension
 
     public X509Certificate2Collection? ServerCertificates { get; private set; }
 
-    public bool DisableRetryExecutionStrategy { get; private set; }
-
     private DbContextOptionsExtensionInfo? _info;
 
     public YdbOptionsExtension()
@@ -25,13 +23,11 @@ private YdbOptionsExtension(YdbOptionsExtension copyFrom) : base(copyFrom)
     {
         CredentialsProvider = copyFrom.CredentialsProvider;
         ServerCertificates = copyFrom.ServerCertificates;
-        DisableRetryExecutionStrategy = copyFrom.DisableRetryExecutionStrategy;
     }
 
     protected override RelationalOptionsExtension Clone() => new YdbOptionsExtension(this);
 
-    public override void ApplyServices(IServiceCollection services) =>
-        services.AddEntityFrameworkYdb(!DisableRetryExecutionStrategy);
+    public override void ApplyServices(IServiceCollection services) => services.AddEntityFrameworkYdb();
 
     public override DbContextOptionsExtensionInfo Info => _info ??= new ExtensionInfo(this);
 
@@ -53,15 +49,6 @@ public YdbOptionsExtension WithServerCertificates(X509Certificate2Collection? se
         return clone;
     }
 
-    public YdbOptionsExtension DisableRetryOnFailure()
-    {
-        var clone = (YdbOptionsExtension)Clone();
-
-        clone.DisableRetryExecutionStrategy = true;
-
-        return clone;
-    }
-
     private sealed class ExtensionInfo(YdbOptionsExtension extension) : RelationalExtensionInfo(extension)
     {
         public override bool IsDatabaseProvider => true;

src/EFCore.Ydb/src/Infrastructure/YdbDbContextOptionsBuilder.cs

@@ -1,20 +1,29 @@
 using System.Security.Cryptography.X509Certificates;
 using EntityFrameworkCore.Ydb.Infrastructure.Internal;
+using EntityFrameworkCore.Ydb.Storage.Internal;
 using Microsoft.EntityFrameworkCore;
 using Microsoft.EntityFrameworkCore.Infrastructure;
+using Microsoft.EntityFrameworkCore.Storage;
+using Ydb.Sdk.Ado.RetryPolicy;
 using Ydb.Sdk.Auth;
 
 namespace EntityFrameworkCore.Ydb.Infrastructure;
 
-public class YdbDbContextOptionsBuilder(DbContextOptionsBuilder optionsBuilder)
+public sealed class YdbDbContextOptionsBuilder(DbContextOptionsBuilder optionsBuilder)
     : RelationalDbContextOptionsBuilder<YdbDbContextOptionsBuilder, YdbOptionsExtension>(optionsBuilder)
 {
-    public YdbDbContextOptionsBuilder WithCredentialsProvider(ICredentialsProvider? credentialsProvider) =>
+    public YdbDbContextOptionsBuilder UseCredentialsProvider(ICredentialsProvider? credentialsProvider) =>
         WithOption(optionsBuilder => optionsBuilder.WithCredentialsProvider(credentialsProvider));
 
-    public YdbDbContextOptionsBuilder WithServerCertificates(X509Certificate2Collection? serverCertificates) =>
+    public YdbDbContextOptionsBuilder UseServerCertificates(X509Certificate2Collection? serverCertificates) =>
         WithOption(optionsBuilder => optionsBuilder.WithServerCertificates(serverCertificates));
 
+    public YdbDbContextOptionsBuilder EnableRetryIdempotence()
+        => UseRetryPolicy(new YdbRetryPolicyConfig { EnableRetryIdempotence = true });
+
+    public YdbDbContextOptionsBuilder UseRetryPolicy(YdbRetryPolicyConfig retryPolicyConfig)
+        => ExecutionStrategy(d => new YdbExecutionStrategy(d, retryPolicyConfig));
+
     public YdbDbContextOptionsBuilder DisableRetryOnFailure() =>
-        WithOption(optionsBuilder => optionsBuilder.DisableRetryOnFailure());
+        ExecutionStrategy(d => new NonRetryingExecutionStrategy(d));
 }

src/EFCore.Ydb/src/Storage/Internal/YdbExecutionStrategy.cs

@@ -1,12 +1,38 @@
 using System;
 using Microsoft.EntityFrameworkCore.Storage;
+using Ydb.Sdk;
 using Ydb.Sdk.Ado;
+using Ydb.Sdk.Ado.RetryPolicy;
 
 namespace EntityFrameworkCore.Ydb.Storage.Internal;
 
-public class YdbExecutionStrategy(ExecutionStrategyDependencies dependencies)
-    : ExecutionStrategy(dependencies, maxRetryCount: 10, maxRetryDelay: TimeSpan.FromSeconds(10)) // TODO User settings
+/// <summary>
+/// Retry strategy for YDB.<br/>
+///
+/// <br/>IMPORTANT:
+/// <br/>- The maximum number of attempts and backoff logic are encapsulated in <see cref="YdbRetryPolicy"/>.
+/// The base ExecutionStrategy parameters (maxRetryCount, maxRetryDelay) are not used.
+/// <br/>- This strategy must be invoked in the correct EF Core context/connection (YDB),
+/// so that exception types and ShouldRetryOn semantics match the provider.
+/// <br/>- The base <see cref="ExecutionStrategy"/> is a good place to emit metrics/logs (attempt number, delay, exception type, etc.).
+/// </summary>
+public class YdbExecutionStrategy(ExecutionStrategyDependencies dependencies, YdbRetryPolicyConfig retryPolicyConfig)
+// We pass "placeholders" to the base class:
+// - MaxAttempts and TimeSpan.Zero are not used in the real retry logic.
+// - Actual limits/delays are driven by YdbRetryPolicy.
+    : ExecutionStrategy(dependencies, retryPolicyConfig.MaxAttempts, TimeSpan.Zero /* unused! */)
 {
-    protected override bool ShouldRetryOn(Exception exception)
-        => exception is YdbException { IsTransient: true };
+    private readonly YdbRetryPolicy _retryPolicy = new(retryPolicyConfig);
+
+    public override bool RetriesOnFailure => true;
+
+    protected override bool ShouldRetryOn(Exception exception) =>
+        exception is YdbException ydbException &&
+        (ydbException.IsTransient || (retryPolicyConfig.EnableRetryIdempotence && ydbException.Code is
+            StatusCode.ClientTransportUnknown or
+            StatusCode.ClientTransportUnavailable or
+            StatusCode.Undetermined));
+
+    protected override TimeSpan? GetNextDelay(Exception lastException) =>
+        _retryPolicy.GetNextDelay((YdbException)lastException, ExceptionsEncountered.Count - 1);
 }

src/EFCore.Ydb/src/Storage/Internal/YdbExecutionStrategyFactory.cs

@@ -1,10 +1,11 @@
 using Microsoft.EntityFrameworkCore.Storage;
+using Ydb.Sdk.Ado.RetryPolicy;
 
 namespace EntityFrameworkCore.Ydb.Storage.Internal;
 
 public class YdbExecutionStrategyFactory(ExecutionStrategyDependencies dependencies)
     : RelationalExecutionStrategyFactory(dependencies)
 {
-    protected override IExecutionStrategy CreateDefaultStrategy(ExecutionStrategyDependencies dependencies)
-        => new YdbExecutionStrategy(dependencies);
+    protected override IExecutionStrategy CreateDefaultStrategy(ExecutionStrategyDependencies dependencies) =>
+        new YdbExecutionStrategy(dependencies, YdbRetryPolicyConfig.Default);
 }

src/EFCore.Ydb/src/Storage/Internal/YdbRelationalConnection.cs

@@ -13,7 +13,6 @@ public class YdbRelationalConnection : RelationalConnection, IYdbRelationalConne
 {
     private readonly ICredentialsProvider? _credentialsProvider;
     private readonly X509Certificate2Collection? _serverCertificates;
-    private readonly bool _disableRetryExecuteStrategy;
 
     public YdbRelationalConnection(RelationalConnectionDependencies dependencies) : base(dependencies)
     {
@@ -22,7 +21,6 @@ public YdbRelationalConnection(RelationalConnectionDependencies dependencies) :
 
         _credentialsProvider = ydbOptionsExtension.CredentialsProvider;
         _serverCertificates = ydbOptionsExtension.ServerCertificates;
-        _disableRetryExecuteStrategy = ydbOptionsExtension.DisableRetryExecutionStrategy;
     }
 
     protected override DbConnection CreateDbConnection()
@@ -40,12 +38,8 @@ public IYdbRelationalConnection Clone()
     {
         var options = new DbContextOptionsBuilder().UseYdb(GetValidatedConnectionString(), builder =>
         {
-            builder.WithCredentialsProvider(_credentialsProvider);
-            builder.WithServerCertificates(_serverCertificates);
-            if (_disableRetryExecuteStrategy)
-            {
-                builder.DisableRetryOnFailure();
-            }
+            builder.UseCredentialsProvider(_credentialsProvider);
+            builder.UseServerCertificates(_serverCertificates);
         }).Options;
 
         return new YdbRelationalConnection(Dependencies with { ContextOptions = options });

src/EFCore.Ydb/test/EntityFrameworkCore.Ydb.FunctionalTests/QueryExpressionInterceptionYdbTestBase.cs

@@ -7,6 +7,7 @@
 using Microsoft.EntityFrameworkCore.TestUtilities;
 using Microsoft.Extensions.DependencyInjection;
 using Xunit;
+using Ydb.Sdk.Ado.RetryPolicy;
 
 namespace EntityFrameworkCore.Ydb.FunctionalTests;
 
@@ -28,7 +29,7 @@ public override DbContextOptionsBuilder AddOptions(DbContextOptionsBuilder build
         {
             new YdbDbContextOptionsBuilder(base.AddOptions(builder))
 #pragma warning disable EF1001
-                .ExecutionStrategy(d => new YdbExecutionStrategy(d));
+                .ExecutionStrategy(d => new YdbExecutionStrategy(d, YdbRetryPolicyConfig.Default));
 #pragma warning restore EF1001
             return builder;
         }

src/Ydb.Sdk/src/Ado/RetryPolicy/YdbRetryPolicy.cs

@@ -85,7 +85,7 @@ internal YdbRetryPolicy(YdbRetryPolicyConfig config, IRandom random) : this(conf
 
         return ydbException.Code switch
         {
-            StatusCode.BadSession or StatusCode.SessionBusy => TimeSpan.Zero,
+            StatusCode.BadSession or StatusCode.SessionBusy or StatusCode.SessionExpired => TimeSpan.Zero,
             StatusCode.Aborted or StatusCode.Undetermined =>
                 FullJitter(_fastBackoffBaseMs, _fastCapBackoffMs, _fastCeiling, attempt, _random),
             StatusCode.Unavailable or StatusCode.ClientTransportUnknown or StatusCode.ClientTransportUnavailable =>