fix leak of PartitionChooserActor (ydb-platform#4132)

shnikd · Apr 26, 2024 · 541d23a · 541d23a
1 parent 5fe6486
commit 541d23a
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 6 deletions.
diff --git a/ydb/core/persqueue/writer/partition_chooser_impl__abstract_chooser_actor.h b/ydb/core/persqueue/writer/partition_chooser_impl__abstract_chooser_actor.h
@@ -58,7 +58,7 @@ class TAbstractPartitionChooserActor: public TActorBootstrapped<TDerived> {
         if (TableHelper.Initialize(ctx, SourceId)) {
             return true;
         }
-        StartIdle();        
+        StartIdle();
         TThis::ReplyError(ErrorCode::BAD_REQUEST, "Bad SourceId", ctx);
         return false;
     }
@@ -67,6 +67,7 @@ class TAbstractPartitionChooserActor: public TActorBootstrapped<TDerived> {
         auto ctx = TActivationContext::ActorContextFor(SelfId());
         TableHelper.CloseKqpSession(ctx);
         PartitionHelper.Close(ctx);
+        TActorBootstrapped<TDerived>::PassAway();
     }
 
     bool NeedTable(const NActors::TActorContext& ctx) {
@@ -78,7 +79,7 @@ class TAbstractPartitionChooserActor: public TActorBootstrapped<TDerived> {
     void InitTable(const NActors::TActorContext& ctx) {
         TThis::Become(&TThis::StateInitTable);
         const auto& pqConfig = AppData(ctx)->PQConfig;
-        TRACE("InitTable: SourceId="<< SourceId 
+        TRACE("InitTable: SourceId="<< SourceId
               << " TopicsAreFirstClassCitizen=" << pqConfig.GetTopicsAreFirstClassCitizen()
               << " UseSrcIdMetaMappingInFirstClass=" <<pqConfig.GetUseSrcIdMetaMappingInFirstClass());
         if (SourceId && pqConfig.GetTopicsAreFirstClassCitizen() && pqConfig.GetUseSrcIdMetaMappingInFirstClass()) {
@@ -259,14 +260,14 @@ class TAbstractPartitionChooserActor: public TActorBootstrapped<TDerived> {
 
 protected:
     void StartIdle() {
-        TThis::Become(&TThis::StateIdle);        
+        TThis::Become(&TThis::StateIdle);
         DEBUG("Start idle");
     }
 
     void HandleIdle(TEvPartitionChooser::TEvRefreshRequest::TPtr&, const TActorContext& ctx) {
         if (PartitionPersisted) {
             SendUpdateRequests(ctx);
-        }        
+        }
     }
 
     STATEFN(StateIdle)  {
@@ -303,7 +304,7 @@ class TAbstractPartitionChooserActor: public TActorBootstrapped<TDerived> {
 
         TThis::Die(ctx);
     }
-    
+
 
 protected:
     const TActorId Parent;

diff --git a/ydb/services/deprecated/persqueue_v0/grpc_pq_write_actor.cpp b/ydb/services/deprecated/persqueue_v0/grpc_pq_write_actor.cpp
@@ -451,6 +451,10 @@ void TWriteSessionActor::Handle(TEvTicketParser::TEvAuthorizeTicketResult::TPtr&
 void TWriteSessionActor::DiscoverPartition(const NActors::TActorContext& ctx) {
     State = ES_WAIT_PARTITION;
 
+    if (PartitionChooser) {
+        ctx.Send(PartitionChooser,  new TEvents::TEvPoison());
+    }
+
     std::optional<ui32> preferedPartition = PreferedPartition == Max<ui32>() ? std::nullopt : std::optional(PreferedPartition);
     PartitionChooser = ctx.RegisterWithSameMailbox(NPQ::CreatePartitionChooserActor(ctx.SelfID, Config, FullConverter, SourceId, preferedPartition));
 }
@@ -894,7 +898,7 @@ void TWriteSessionActor::HandleWakeup(const TActorContext& ctx) {
         ctx.Send(PartitionChooser, new NPQ::TEvPartitionChooser::TEvRefreshRequest());
         LastSourceIdUpdate = now + SOURCEID_UPDATE_PERIOD;
     }
-    
+
     if (now >= LogSessionDeadline) {
         LogSession(ctx);
     }

diff --git a/ydb/services/persqueue_v1/actors/write_session_actor.cpp b/ydb/services/persqueue_v1/actors/write_session_actor.cpp
@@ -627,6 +627,10 @@ template<bool UseMigrationProtocol>
 void TWriteSessionActor<UseMigrationProtocol>::DiscoverPartition(const NActors::TActorContext& ctx) {
     State = ES_WAIT_PARTITION;
 
+    if (PartitionChooser) {
+        ctx.Send(PartitionChooser,  new TEvents::TEvPoison());
+    }
+
     std::optional<ui32> preferedPartition = PreferedPartition == Max<ui32>() ? std::nullopt : std::optional(PreferedPartition);
     PartitionChooser = ctx.RegisterWithSameMailbox(NPQ::CreatePartitionChooserActor(ctx.SelfID, Config, FullConverter, SourceId, preferedPartition));
 }