Files
vigilcare-clinical/VigilCareClinicalAPI/DataLake/DataLakeWriterService.cs
T

264 lines
10 KiB
C#

using System.Text.Json;
using Confluent.Kafka;
using Microsoft.Extensions.Options;
using Minio.DataModel.Args;
public sealed class DataLakeWriterService : BackgroundService
{
private readonly KafkaOptions _kafkaOptions;
private readonly DataLakeOptions _opts;
private readonly MinioOptions _minioOpts;
private readonly ILogger<DataLakeWriterService> _logger;
// Buffer key: identifies one Parquet file-to-be.
// Events sharing a topic, date, and Kafka partition land in the same file.
private record BufferKey(string Topic, string DatePath, int Partition);
private record BufferedEvent(string Payload, long Offset);
private readonly Dictionary<BufferKey, List<BufferedEvent>> _buffer = new();
// Track the highest offset per topic-partition for post-flush commit.
private readonly Dictionary<TopicPartition, TopicPartitionOffset> _highWatermarks = new();
public DataLakeWriterService(
IOptions<KafkaOptions> kafkaOptions,
IOptions<DataLakeOptions> opts,
IOptions<MinioOptions> minioOpts,
ILogger<DataLakeWriterService> logger)
{
_kafkaOptions = kafkaOptions.Value;
_opts = opts.Value;
_minioOpts = minioOpts.Value;
_logger = logger;
}
protected override async Task ExecuteAsync(CancellationToken ct)
{
var consumerConfig = new ConsumerConfig
{
BootstrapServers = _kafkaOptions.BootstrapServers,
GroupId = "data-lake-writer",
AutoOffsetReset = AutoOffsetReset.Earliest,
EnableAutoCommit = false,
};
using var consumer = new ConsumerBuilder<string, string>(consumerConfig).Build();
consumer.Subscribe(new[]
{
_kafkaOptions.Topics.ObservationRecorded,
_kafkaOptions.Topics.AlertGenerated,
_kafkaOptions.Topics.EncounterStatusChanged,
});
_logger.LogInformation(
"DataLakeWriterService started. FlushCount={FlushCount} FlushIntervalSeconds={FlushInterval}",
_opts.FlushCount, _opts.FlushIntervalSeconds);
var lastFlush = DateTimeOffset.UtcNow;
try
{
while (!ct.IsCancellationRequested)
{
ConsumeResult<string, string>? result;
try
{
result = consumer.Consume(TimeSpan.FromMilliseconds(500));
}
catch (ConsumeException ex)
{
_logger.LogError(ex, "DataLakeWriter consume error");
continue;
}
if (result is not null)
AddToBuffer(result);
var totalBuffered = _buffer.Values.Sum(v => v.Count);
var shouldFlushCount = totalBuffered >= _opts.FlushCount;
var shouldFlushTime = DateTimeOffset.UtcNow - lastFlush
>= TimeSpan.FromSeconds(_opts.FlushIntervalSeconds);
if ((shouldFlushCount || shouldFlushTime) && totalBuffered > 0)
{
await FlushAsync(consumer, ct);
lastFlush = DateTimeOffset.UtcNow;
}
}
}
finally
{
// Final flush on shutdown so buffered events are not lost.
// Use CancellationToken.None — host shutdown cancels the execute token before
// MinIO uploads finish, which surfaces as TaskCanceledException noise.
if (_buffer.Values.Sum(v => v.Count) > 0)
{
try { await FlushAsync(consumer, CancellationToken.None); }
catch (Exception ex)
{
_logger.LogError(ex, "DataLakeWriter shutdown flush failed — some events may be re-read on next start");
}
}
consumer.Close();
}
}
private void AddToBuffer(ConsumeResult<string, string> result)
{
var datePath = ExtractDatePath(result.Topic, result.Message.Value);
var key = new BufferKey(result.Topic, datePath, result.Partition.Value);
if (!_buffer.TryGetValue(key, out var list))
{
list = new List<BufferedEvent>();
_buffer[key] = list;
}
list.Add(new BufferedEvent(result.Message.Value, result.Offset.Value));
// Track highest offset per topic-partition for post-flush commit.
var tp = new TopicPartition(result.Topic, result.Partition);
_highWatermarks[tp] = new TopicPartitionOffset(tp, result.Offset + 1);
}
private async Task FlushAsync(IConsumer<string, string> consumer, CancellationToken ct)
{
var flushedKeys = new HashSet<BufferKey>();
var failedPartitions = new HashSet<int>();
foreach (var (key, events) in _buffer)
{
if (events.Count == 0) continue;
try
{
var firstOffset = events.Min(e => e.Offset);
var objectKey = BuildObjectKey(key, firstOffset);
var bytes = await BuildParquetAsync(key.Topic, events, key.Partition);
await UploadToMinioAsync(objectKey, bytes, ct);
flushedKeys.Add(key);
_logger.LogInformation(
"[DATA-LAKE] Wrote {Count} events → {ObjectKey} ({Bytes} bytes)",
events.Count, objectKey, bytes.Length);
}
catch (Exception ex)
{
failedPartitions.Add(key.Partition);
if (ex is OperationCanceledException && ct.IsCancellationRequested)
{
_logger.LogInformation(
"[DATA-LAKE] Flush canceled for key {Key} during shutdown", key);
}
else
{
_logger.LogError(ex, "[DATA-LAKE] Failed to write file for key {Key}", key);
}
}
}
// Only commit offsets for topic-partitions that had no failures.
var safeOffsets = _highWatermarks
.Where(kv => !failedPartitions.Contains(kv.Key.Partition))
.Select(kv => kv.Value)
.ToList();
if (safeOffsets.Count > 0)
{
consumer.Commit(safeOffsets);
_logger.LogInformation(
"[DATA-LAKE] Committed offsets for {PartitionCount} partitions after flushing {FileCount} files",
safeOffsets.Count, flushedKeys.Count);
}
if (failedPartitions.Count > 0)
{
_logger.LogWarning(
"[DATA-LAKE] Retained buffers for {FailedCount} failed partitions — will retry next flush",
failedPartitions.Count);
}
// Clear only successfully flushed keys; retain failed ones for retry.
foreach (var key in flushedKeys)
_buffer.Remove(key);
// Clear watermarks only for partitions with no failures.
foreach (var tp in _highWatermarks.Keys.ToList())
{
if (!failedPartitions.Contains(tp.Partition))
_highWatermarks.Remove(tp);
}
}
private string ExtractDatePath(string topic, string payload) =>
DataLakeEventParser.ExtractDatePath(topic, payload, _kafkaOptions.Topics);
private async Task<byte[]> BuildParquetAsync(
string topic, List<BufferedEvent> events, int partition)
{
if (topic == _kafkaOptions.Topics.ObservationRecorded)
{
var rows = events.Select(e => DataLakeEventParser.ParseObservationRow(
e.Payload, e.Offset, partition)).ToList();
return await ParquetFileBuilder.BuildObservationsAsync(rows);
}
if (topic == _kafkaOptions.Topics.AlertGenerated)
{
var rows = events.Select(e => DataLakeEventParser.ParseAlertRow(
e.Payload, e.Offset, partition)).ToList();
return await ParquetFileBuilder.BuildAlertsAsync(rows);
}
if (topic == _kafkaOptions.Topics.EncounterStatusChanged)
{
var rows = events.Select(e => DataLakeEventParser.ParseEncounterRow(
e.Payload, e.Offset, partition)).ToList();
return await ParquetFileBuilder.BuildEncountersAsync(rows);
}
throw new InvalidOperationException($"Unknown topic: {topic}");
}
// ---------------------------------------------------------------------------
// Object key and date partition helpers
// ---------------------------------------------------------------------------
// File path: observations/2025/01/15/partition-0-offset-0000001000.parquet
// The date comes from the event timestamp, not the wall clock.
// Events from the same encounter that cross midnight are written into the date
// bucket matching their recorded_at timestamp — consistent with how Athena and
// Spark partition-prune by event time, not ingest time.
private string BuildObjectKey(BufferKey key, long firstOffset)
{
var folder = key.Topic switch
{
var t when t == _kafkaOptions.Topics.ObservationRecorded => "observations",
var t when t == _kafkaOptions.Topics.AlertGenerated => "alerts",
var t when t == _kafkaOptions.Topics.EncounterStatusChanged => "encounters",
_ => "unknown",
};
return $"{folder}/{key.DatePath}/partition-{key.Partition}-offset-{firstOffset:D10}.parquet";
}
// ---------------------------------------------------------------------------
// MinIO upload
// ---------------------------------------------------------------------------
private async Task UploadToMinioAsync(string objectKey, byte[] bytes, CancellationToken ct)
{
var client = MinioClientFactory.Build(_minioOpts);
var bucket = _opts.BucketName;
var exists = await client.BucketExistsAsync(
new BucketExistsArgs().WithBucket(bucket), ct);
if (!exists)
await client.MakeBucketAsync(new MakeBucketArgs().WithBucket(bucket), ct);
await client.PutObjectAsync(new PutObjectArgs()
.WithBucket(bucket)
.WithObject(objectKey)
.WithStreamData(new MemoryStream(bytes))
.WithObjectSize(bytes.Length)
.WithContentType("application/octet-stream"),
ct);
}
}