Fix leak of PartitionChooserActor (#4131)

Author: nshestakov

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,8 @@ 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 +80,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 +261,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 +305,7 @@ class TAbstractPartitionChooserActor: public TActorBootstrapped<TDerived> {
 
         TThis::Die(ctx);
     }
-    
+
 
 protected:
     const TActorId Parent;

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);
     }

ydb/services/persqueue_v1/actors/write_session_actor.ipp

@@ -624,6 +624,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));
 }