Skip to content

Commit 931f0c6

Browse files
kdenneykdenney
authored
[PM-28265] storage reconciliation job (#6615)
1 parent 9573cab commit 931f0c6

File tree

10 files changed

+993
-15
lines changed

10 files changed

+993
-15
lines changed

‎src/Billing/Controllers/JobsController.cs‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
using Bit.Billing.Jobs;
2+
using Bit.Core.Utilities;
3+
using Microsoft.AspNetCore.Mvc;
4+
5+
namespace Bit.Billing.Controllers;
6+
7+
[Route("jobs")]
8+
[SelfHosted(NotSelfHostedOnly = true)]
9+
[RequireLowerEnvironment]
10+
public class JobsController(
11+
JobsHostedService jobsHostedService) : Controller
12+
{
13+
[HttpPost("run/{jobName}")]
14+
public async Task<IActionResult> RunJobAsync(string jobName)
15+
{
16+
if (jobName == nameof(ReconcileAdditionalStorageJob))
17+
{
18+
await jobsHostedService.RunJobAdHocAsync<ReconcileAdditionalStorageJob>();
19+
return Ok(new { message = $"Job {jobName} scheduled successfully" });
20+
}
21+
22+
return BadRequest(new { error = $"Unknown job name: {jobName}" });
23+
}
24+
25+
[HttpPost("stop/{jobName}")]
26+
public async Task<IActionResult> StopJobAsync(string jobName)
27+
{
28+
if (jobName == nameof(ReconcileAdditionalStorageJob))
29+
{
30+
await jobsHostedService.InterruptAdHocJobAsync<ReconcileAdditionalStorageJob>();
31+
return Ok(new { message = $"Job {jobName} queued for cancellation" });
32+
}
33+
34+
return BadRequest(new { error = $"Unknown job name: {jobName}" });
35+
}
36+
}

‎src/Billing/Jobs/AliveJob.cs‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,4 +10,13 @@ protected override Task ExecuteJobAsync(IJobExecutionContext context)
1010
_logger.LogInformation(Core.Constants.BypassFiltersEventId, null, "Billing service is alive!");
1111
return Task.FromResult(0);
1212
}
13+
14+
public static ITrigger GetTrigger()
15+
{
16+
return TriggerBuilder.Create()
17+
.WithIdentity("EveryTopOfTheHourTrigger")
18+
.StartNow()
19+
.WithCronSchedule("0 0 * * * ?")
20+
.Build();
21+
}
1322
}

‎src/Billing/Jobs/JobsHostedService.cs‎

Lines changed: 62 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,29 +1,27 @@
1-
using Bit.Core.Jobs;
1+
using Bit.Core.Exceptions;
2+
using Bit.Core.Jobs;
23
using Bit.Core.Settings;
34
using Quartz;
45

56
namespace Bit.Billing.Jobs;
67

7-
public class JobsHostedService : BaseJobsHostedService
8+
public class JobsHostedService(
9+
GlobalSettings globalSettings,
10+
IServiceProvider serviceProvider,
11+
ILogger<JobsHostedService> logger,
12+
ILogger<JobListener> listenerLogger,
13+
ISchedulerFactory schedulerFactory)
14+
: BaseJobsHostedService(globalSettings, serviceProvider, logger, listenerLogger)
815
{
9-
public JobsHostedService(
10-
GlobalSettings globalSettings,
11-
IServiceProvider serviceProvider,
12-
ILogger<JobsHostedService> logger,
13-
ILogger<JobListener> listenerLogger)
14-
: base(globalSettings, serviceProvider, logger, listenerLogger) { }
16+
private List<JobKey> AdHocJobKeys { get; } = [];
17+
private IScheduler? _adHocScheduler;
1518

1619
public override async Task StartAsync(CancellationToken cancellationToken)
1720
{
18-
var everyTopOfTheHourTrigger = TriggerBuilder.Create()
19-
.WithIdentity("EveryTopOfTheHourTrigger")
20-
.StartNow()
21-
.WithCronSchedule("0 0 * * * ?")
22-
.Build();
23-
2421
Jobs = new List<Tuple<Type, ITrigger>>
2522
{
26-
new Tuple<Type, ITrigger>(typeof(AliveJob), everyTopOfTheHourTrigger)
23+
new(typeof(AliveJob), AliveJob.GetTrigger()),
24+
new(typeof(ReconcileAdditionalStorageJob), ReconcileAdditionalStorageJob.GetTrigger())
2725
};
2826

2927
await base.StartAsync(cancellationToken);
@@ -33,5 +31,54 @@ public static void AddJobsServices(IServiceCollection services)
3331
{
3432
services.AddTransient<AliveJob>();
3533
services.AddTransient<SubscriptionCancellationJob>();
34+
services.AddTransient<ReconcileAdditionalStorageJob>();
35+
// add this service as a singleton so we can inject it where needed
36+
services.AddSingleton<JobsHostedService>();
37+
services.AddHostedService(sp => sp.GetRequiredService<JobsHostedService>());
38+
}
39+
40+
public async Task InterruptAdHocJobAsync<T>(CancellationToken cancellationToken = default) where T : class, IJob
41+
{
42+
if (_adHocScheduler == null)
43+
{
44+
throw new InvalidOperationException("AdHocScheduler is null, cannot interrupt ad-hoc job.");
45+
}
46+
47+
var jobKey = AdHocJobKeys.FirstOrDefault(j => j.Name == typeof(T).ToString());
48+
if (jobKey == null)
49+
{
50+
throw new NotFoundException($"Cannot find job key: {typeof(T)}, not running?");
51+
}
52+
logger.LogInformation("CANCELLING ad-hoc job with key: {JobKey}", jobKey);
53+
AdHocJobKeys.Remove(jobKey);
54+
await _adHocScheduler.Interrupt(jobKey, cancellationToken);
55+
}
56+
57+
public async Task RunJobAdHocAsync<T>(CancellationToken cancellationToken = default) where T : class, IJob
58+
{
59+
_adHocScheduler ??= await schedulerFactory.GetScheduler(cancellationToken);
60+
61+
var jobKey = new JobKey(typeof(T).ToString());
62+
63+
var currentlyExecuting = await _adHocScheduler.GetCurrentlyExecutingJobs(cancellationToken);
64+
if (currentlyExecuting.Any(j => j.JobDetail.Key.Equals(jobKey)))
65+
{
66+
throw new InvalidOperationException($"Job {jobKey} is already running");
67+
}
68+
69+
AdHocJobKeys.Add(jobKey);
70+
71+
var job = JobBuilder.Create<T>()
72+
.WithIdentity(jobKey)
73+
.Build();
74+
75+
var trigger = TriggerBuilder.Create()
76+
.WithIdentity(typeof(T).ToString())
77+
.StartNow()
78+
.Build();
79+
80+
logger.LogInformation("Scheduling ad-hoc job with key: {JobKey}", jobKey);
81+
82+
await _adHocScheduler.ScheduleJob(job, trigger, cancellationToken);
3683
}
3784
}

‎src/Billing/Jobs/ReconcileAdditionalStorageJob.cs‎

Lines changed: 207 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,207 @@
1+
using System.Globalization;
2+
using System.Text.Json;
3+
using Bit.Billing.Services;
4+
using Bit.Core;
5+
using Bit.Core.Billing.Constants;
6+
using Bit.Core.Jobs;
7+
using Bit.Core.Services;
8+
using Quartz;
9+
using Stripe;
10+
11+
namespace Bit.Billing.Jobs;
12+
13+
public class ReconcileAdditionalStorageJob(
14+
IStripeFacade stripeFacade,
15+
ILogger<ReconcileAdditionalStorageJob> logger,
16+
IFeatureService featureService) : BaseJob(logger)
17+
{
18+
private const string _storageGbMonthlyPriceId = "storage-gb-monthly";
19+
private const string _storageGbAnnuallyPriceId = "storage-gb-annually";
20+
private const string _personalStorageGbAnnuallyPriceId = "personal-storage-gb-annually";
21+
private const int _storageGbToRemove = 4;
22+
23+
protected override async Task ExecuteJobAsync(IJobExecutionContext context)
24+
{
25+
if (!featureService.IsEnabled(FeatureFlagKeys.PM28265_EnableReconcileAdditionalStorageJob))
26+
{
27+
logger.LogInformation("Skipping ReconcileAdditionalStorageJob, feature flag off.");
28+
return;
29+
}
30+
31+
var liveMode = featureService.IsEnabled(FeatureFlagKeys.PM28265_ReconcileAdditionalStorageJobEnableLiveMode);
32+
33+
// Execution tracking
34+
var subscriptionsFound = 0;
35+
var subscriptionsUpdated = 0;
36+
var subscriptionsWithErrors = 0;
37+
var failures = new List<string>();
38+
39+
logger.LogInformation("Starting ReconcileAdditionalStorageJob (live mode: {LiveMode})", liveMode);
40+
41+
var priceIds = new[] { _storageGbMonthlyPriceId, _storageGbAnnuallyPriceId, _personalStorageGbAnnuallyPriceId };
42+
43+
foreach (var priceId in priceIds)
44+
{
45+
var options = new SubscriptionListOptions
46+
{
47+
Limit = 100,
48+
Status = StripeConstants.SubscriptionStatus.Active,
49+
Price = priceId
50+
};
51+
52+
await foreach (var subscription in stripeFacade.ListSubscriptionsAutoPagingAsync(options))
53+
{
54+
if (context.CancellationToken.IsCancellationRequested)
55+
{
56+
logger.LogWarning(
57+
"Job cancelled!! Exiting. Progress at time of cancellation: Subscriptions found: {SubscriptionsFound}, " +
58+
"Updated: {SubscriptionsUpdated}, Errors: {SubscriptionsWithErrors}{Failures}",
59+
subscriptionsFound,
60+
liveMode
61+
? subscriptionsUpdated
62+
: $"(In live mode, would have updated) {subscriptionsUpdated}",
63+
subscriptionsWithErrors,
64+
failures.Count > 0
65+
? $", Failures: {Environment.NewLine}{string.Join(Environment.NewLine, failures)}"
66+
: string.Empty
67+
);
68+
return;
69+
}
70+
71+
if (subscription == null)
72+
{
73+
continue;
74+
}
75+
76+
logger.LogInformation("Processing subscription: {SubscriptionId}", subscription.Id);
77+
subscriptionsFound++;
78+
79+
if (subscription.Metadata?.TryGetValue(StripeConstants.MetadataKeys.StorageReconciled2025, out var dateString) == true)
80+
{
81+
if (DateTime.TryParse(dateString, null, DateTimeStyles.RoundtripKind, out var dateProcessed))
82+
{
83+
logger.LogInformation("Skipping subscription {SubscriptionId} - already processed on {Date}",
84+
subscription.Id,
85+
dateProcessed.ToString("f"));
86+
continue;
87+
}
88+
}
89+
90+
var updateOptions = BuildSubscriptionUpdateOptions(subscription, priceId);
91+
92+
if (updateOptions == null)
93+
{
94+
logger.LogInformation("Skipping subscription {SubscriptionId} - no updates needed", subscription.Id);
95+
continue;
96+
}
97+
98+
subscriptionsUpdated++;
99+
100+
if (!liveMode)
101+
{
102+
logger.LogInformation(
103+
"Not live mode (dry-run): Would have updated subscription {SubscriptionId} with item changes: {NewLine}{UpdateOptions}",
104+
subscription.Id,
105+
Environment.NewLine,
106+
JsonSerializer.Serialize(updateOptions));
107+
continue;
108+
}
109+
110+
try
111+
{
112+
await stripeFacade.UpdateSubscription(subscription.Id, updateOptions);
113+
logger.LogInformation("Successfully updated subscription: {SubscriptionId}", subscription.Id);
114+
}
115+
catch (Exception ex)
116+
{
117+
subscriptionsWithErrors++;
118+
failures.Add($"Subscription {subscription.Id}: {ex.Message}");
119+
logger.LogError(ex, "Failed to update subscription {SubscriptionId}: {ErrorMessage}",
120+
subscription.Id, ex.Message);
121+
}
122+
}
123+
}
124+
125+
logger.LogInformation(
126+
"ReconcileAdditionalStorageJob completed. Subscriptions found: {SubscriptionsFound}, " +
127+
"Updated: {SubscriptionsUpdated}, Errors: {SubscriptionsWithErrors}{Failures}",
128+
subscriptionsFound,
129+
liveMode
130+
? subscriptionsUpdated
131+
: $"(In live mode, would have updated) {subscriptionsUpdated}",
132+
subscriptionsWithErrors,
133+
failures.Count > 0
134+
? $", Failures: {Environment.NewLine}{string.Join(Environment.NewLine, failures)}"
135+
: string.Empty
136+
);
137+
}
138+
139+
private SubscriptionUpdateOptions? BuildSubscriptionUpdateOptions(
140+
Subscription subscription,
141+
string targetPriceId)
142+
{
143+
if (subscription.Items?.Data == null)
144+
{
145+
return null;
146+
}
147+
148+
var updateOptions = new SubscriptionUpdateOptions
149+
{
150+
ProrationBehavior = StripeConstants.ProrationBehavior.CreateProrations,
151+
Metadata = new Dictionary<string, string>
152+
{
153+
[StripeConstants.MetadataKeys.StorageReconciled2025] = DateTime.UtcNow.ToString("o")
154+
},
155+
Items = []
156+
};
157+
158+
var hasUpdates = false;
159+
160+
foreach (var item in subscription.Items.Data.Where(item => item?.Price?.Id == targetPriceId))
161+
{
162+
hasUpdates = true;
163+
var currentQuantity = item.Quantity;
164+
165+
if (currentQuantity > _storageGbToRemove)
166+
{
167+
var newQuantity = currentQuantity - _storageGbToRemove;
168+
logger.LogInformation(
169+
"Subscription {SubscriptionId}: reducing quantity from {CurrentQuantity} to {NewQuantity} for price {PriceId}",
170+
subscription.Id,
171+
currentQuantity,
172+
newQuantity,
173+
item.Price.Id);
174+
175+
updateOptions.Items.Add(new SubscriptionItemOptions
176+
{
177+
Id = item.Id,
178+
Quantity = newQuantity
179+
});
180+
}
181+
else
182+
{
183+
logger.LogInformation("Subscription {SubscriptionId}: deleting storage item with quantity {CurrentQuantity} for price {PriceId}",
184+
subscription.Id,
185+
currentQuantity,
186+
item.Price.Id);
187+
188+
updateOptions.Items.Add(new SubscriptionItemOptions
189+
{
190+
Id = item.Id,
191+
Deleted = true
192+
});
193+
}
194+
}
195+
196+
return hasUpdates ? updateOptions : null;
197+
}
198+
199+
public static ITrigger GetTrigger()
200+
{
201+
return TriggerBuilder.Create()
202+
.WithIdentity("EveryMorningTrigger")
203+
.StartNow()
204+
.WithCronSchedule("0 0 16 * * ?") // 10am CST daily; the pods execute in UTC time
205+
.Build();
206+
}
207+
}

‎src/Billing/Services/IStripeFacade.cs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,11 @@ Task<StripeList<Subscription>> ListSubscriptions(
7878
RequestOptions requestOptions = null,
7979
CancellationToken cancellationToken = default);
8080

81+
IAsyncEnumerable<Subscription> ListSubscriptionsAutoPagingAsync(
82+
SubscriptionListOptions options = null,
83+
RequestOptions requestOptions = null,
84+
CancellationToken cancellationToken = default);
85+
8186
Task<Subscription> GetSubscription(
8287
string subscriptionId,
8388
SubscriptionGetOptions subscriptionGetOptions = null,

‎src/Billing/Services/Implementations/StripeFacade.cs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,12 @@ public async Task<StripeList<Subscription>> ListSubscriptions(SubscriptionListOp
9898
CancellationToken cancellationToken = default) =>
9999
await _subscriptionService.ListAsync(options, requestOptions, cancellationToken);
100100

101+
public IAsyncEnumerable<Subscription> ListSubscriptionsAutoPagingAsync(
102+
SubscriptionListOptions options = null,
103+
RequestOptions requestOptions = null,
104+
CancellationToken cancellationToken = default) =>
105+
_subscriptionService.ListAutoPagingAsync(options, requestOptions, cancellationToken);
106+
101107
public async Task<Subscription> GetSubscription(
102108
string subscriptionId,
103109
SubscriptionGetOptions subscriptionGetOptions = null,

‎src/Core/Billing/Constants/StripeConstants.cs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@ public static class MetadataKeys
6565
public const string Region = "region";
6666
public const string RetiredBraintreeCustomerId = "btCustomerId_old";
6767
public const string UserId = "userId";
68+
public const string StorageReconciled2025 = "storage_reconciled_2025";
6869
}
6970

7071
public static class PaymentBehavior

‎src/Core/Constants.cs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,8 @@ public static class FeatureFlagKeys
197197
public const string PM26793_FetchPremiumPriceFromPricingService = "pm-26793-fetch-premium-price-from-pricing-service";
198198
public const string PM23341_Milestone_2 = "pm-23341-milestone-2";
199199
public const string PM26462_Milestone_3 = "pm-26462-milestone-3";
200+
public const string PM28265_EnableReconcileAdditionalStorageJob = "pm-28265-enable-reconcile-additional-storage-job";
201+
public const string PM28265_ReconcileAdditionalStorageJobEnableLiveMode = "pm-28265-reconcile-additional-storage-job-enable-live-mode";
200202

201203
/* Key Management Team */
202204
public const string ReturnErrorOnExistingKeypair = "return-error-on-existing-keypair";

0 commit comments

Comments
 (0)