Refactoring - optimize parameters of CreatePartitionWriter (#669)

Author: nshestakov

ydb/core/kafka_proxy/actors/kafka_produce_actor.cpp

@@ -132,9 +132,7 @@ void TKafkaProduceActor::HandleInit(TEvTxProxySchemeCache::TEvNavigateKeySetResu
             if (info.SecurityObject->CheckAccess(NACLib::EAccessRights::UpdateRow, *Context->UserToken)) {
                 topic.Status = OK;
                 topic.ExpirationTime = now + TOPIC_OK_EXPIRATION_INTERVAL;
-                for(auto& p : info.PQGroupInfo->Description.GetPartitions()) {
-                    topic.partitions[p.GetPartitionId()] = p.GetTabletId();
-                }
+                topic.PartitionChooser = CreatePartitionChooser(info.PQGroupInfo->Description);
             } else {
                 KAFKA_LOG_W("Produce actor: Unauthorized PRODUCE to topic '" << topicPath << "'");
                 topic.Status = UNAUTHORIZED;
@@ -178,7 +176,7 @@ void TKafkaProduceActor::Handle(TEvTxProxySchemeCache::TEvWatchNotifyDeleted::TP
     auto& topicInfo = Topics[path];
     topicInfo.Status = NOT_FOUND;
     topicInfo.ExpirationTime = ctx.Now() + TOPIC_NOT_FOUND_EXPIRATION_INTERVAL;
-    topicInfo.partitions.clear();
+    topicInfo.PartitionChooser.reset();
 }
 
 void TKafkaProduceActor::Handle(TEvTxProxySchemeCache::TEvWatchNotifyUpdated::TPtr& ev, const TActorContext& ctx) {
@@ -192,10 +190,7 @@ void TKafkaProduceActor::Handle(TEvTxProxySchemeCache::TEvWatchNotifyUpdated::TP
     }
     topic.Status = OK;
     topic.ExpirationTime = ctx.Now() + TOPIC_OK_EXPIRATION_INTERVAL;
-    topic.partitions.clear();
-    for (auto& p : e->Result->GetPathDescription().GetPersQueueGroup().GetPartitions()) {
-        topic.partitions[p.GetPartitionId()] = p.GetTabletId();
-    }
+    topic.PartitionChooser = CreatePartitionChooser(e->Result->GetPathDescription().GetPersQueueGroup());
 }
 
 void TKafkaProduceActor::Handle(TEvKafka::TEvProduceRequest::TPtr request, const TActorContext& ctx) {
@@ -578,17 +573,17 @@ std::pair<TKafkaProduceActor::ETopicStatus, TActorId> TKafkaProduceActor::Partit
         return { OK, writerInfo.ActorId };
     }
 
-    auto& partitions = topicInfo.partitions;
-    auto pit = partitions.find(partitionId);
-    if (pit == partitions.end()) {
+    auto* partition = topicInfo.PartitionChooser->GetPartition(partitionId);
+    if (!partition) {
         return { NOT_FOUND, TActorId{} };
     }
 
-    auto tabletId = pit->second;
     TPartitionWriterOpts opts;
     opts.WithDeduplication(false)
+        .WithSourceId(SourceId)
+        .WithTopicPath(topicPath)
         .WithCheckRequestUnits(topicInfo.MeteringMode, Context->RlContext);
-    auto* writerActor = CreatePartitionWriter(SelfId(), topicPath, tabletId, partitionId, {/*expectedGeneration*/}, SourceId, opts);
+    auto* writerActor = CreatePartitionWriter(SelfId(), partition->TabletId, partitionId, opts);
 
     auto& writerInfo = partitionWriters[partitionId];
     writerInfo.ActorId = ctx.RegisterWithSameMailbox(writerActor);

ydb/core/kafka_proxy/actors/kafka_produce_actor.h

@@ -181,9 +181,7 @@ class TKafkaProduceActor: public NActors::TActorBootstrapped<TKafkaProduceActor>
         TInstant ExpirationTime;
 
         NKikimrPQ::TPQTabletConfig::EMeteringMode MeteringMode;
-
-        // partitioId -> tabletId
-        std::unordered_map<ui32, ui64> partitions;
+        std::shared_ptr<IPartitionChooser> PartitionChooser;
     };
     std::map<TString, TTopicInfo> Topics;
 

ydb/core/persqueue/writer/partition_chooser.h

@@ -48,6 +48,25 @@ struct TEvPartitionChooser {
 
 };
 
+class IPartitionChooser {
+public:
+    struct TPartitionInfo {
+        TPartitionInfo(ui32 partitionId, ui64 tabletId)
+            : PartitionId(partitionId)
+            , TabletId(tabletId) {}
+
+        ui32 PartitionId;
+        ui64 TabletId;
+    };
+
+    virtual ~IPartitionChooser() = default;
+
+    virtual const TPartitionInfo* GetPartition(const TString& sourceId) const = 0;
+    virtual const TPartitionInfo* GetPartition(ui32 partitionId) const = 0;
+};
+
+std::shared_ptr<IPartitionChooser> CreatePartitionChooser(const NKikimrSchemeOp::TPersQueueGroupDescription& config, bool withoutHash = false);
+
 NActors::IActor* CreatePartitionChooserActor(TActorId parentId,
                                              const NKikimrSchemeOp::TPersQueueGroupDescription& config,
                                              NPersQueue::TTopicConverterPtr& fullConverter,