Fix direct write status code (#7735)

Author: nshestakov

ydb/core/persqueue/read_balancer__balancing.cpp

@@ -285,6 +285,8 @@ void TPartitionFamily::AfterRelease() {
     Partitions.clear();
     Partitions.insert(Partitions.end(), RootPartitions.begin(), RootPartitions.end());
 
+    LockedPartitions.clear();
+
     ClassifyPartitions();
     UpdatePartitionMapping(Partitions);
     // After reducing the number of partitions in the family, the list of reading sessions that can read this family may expand.
@@ -464,7 +466,7 @@ bool TPartitionFamily::PossibleForBalance(TSession* session) {
 
 void TPartitionFamily::ClassifyPartitions() {
     auto [activePartitionCount, inactivePartitionCount] = ClassifyPartitions(Partitions);
-    ChangePartitionCounters(activePartitionCount, inactivePartitionCount);
+    ChangePartitionCounters(activePartitionCount - ActivePartitionCount, inactivePartitionCount - InactivePartitionCount);
 }
 
 template<typename TPartitions>
@@ -1792,10 +1794,9 @@ void TBalancer::Handle(TEvPersQueue::TEvGetReadSessionsInfo::TPtr& ev, const TAc
             pi->SetPartition(partitionId);
 
             auto* family = consumer->FindFamily(partitionId);
-            if (family && family->LockedPartitions.contains(partitionId)) {
+            if (family && family->Session && family->LockedPartitions.contains(partitionId)) {
                 auto* session = family->Session;
 
-                Y_ABORT_UNLESS(session != nullptr);
                 pi->SetClientNode(session->ClientNode);
                 pi->SetProxyNodeId(session->ProxyNodeId);
                 pi->SetSession(session->SessionName);

ydb/core/persqueue/writer/partition_chooser_impl__abstract_chooser_actor.h

@@ -217,7 +217,7 @@ class TAbstractPartitionChooserActor: public TActorBootstrapped<TDerived> {
         TThis::Become(&TThis::StateCheckPartition);
 
         if (!Partition) {
-            return ReplyError(ErrorCode::INITIALIZING, "Partition not choosed", ctx);
+            return ReplyError(TThis::PreferedPartition ? ErrorCode::WRITE_ERROR_PARTITION_INACTIVE : ErrorCode::INITIALIZING, "Partition not choosed", ctx);
         }
 
         PartitionHelper.Open(Partition->TabletId, ctx);

ydb/core/persqueue/writer/partition_chooser_impl__sm_chooser_actor.h

@@ -71,7 +71,7 @@ class TSMPartitionChooserActor: public TAbstractPartitionChooserActor<TSMPartiti
 
         const auto* node = Graph.GetPartition(TThis::TableHelper.PartitionId().value());
         if (!node) {
-            // The partition where the writting was performed earlier has already been deleted. 
+            // The partition where the writting was performed earlier has already been deleted.
             // We can write without taking into account the hierarchy of the partition.
             TThis::Partition = BoundaryPartition;
             return OnPartitionChosen(ctx);
@@ -109,7 +109,7 @@ class TSMPartitionChooserActor: public TAbstractPartitionChooserActor<TSMPartiti
     void GetOwnershipFast(const TActorContext &ctx) {
         TThis::Become(&TThis::StateOwnershipFast);
         if (!BoundaryPartition) {
-            return TThis::ReplyError(ErrorCode::INITIALIZING, "A partition not choosed", ctx);
+            return TThis::ReplyError(TThis::PreferedPartition ? ErrorCode::WRITE_ERROR_PARTITION_INACTIVE : ErrorCode::INITIALIZING, "A partition not choosed", ctx);
         }
 
         DEBUG("GetOwnershipFast Partition=" << BoundaryPartition->PartitionId << " TabletId=" << BoundaryPartition->TabletId);
@@ -120,7 +120,7 @@ class TSMPartitionChooserActor: public TAbstractPartitionChooserActor<TSMPartiti
 
     void HandleFast(NKikimr::TEvPQ::TEvCheckPartitionStatusResponse::TPtr& ev, const NActors::TActorContext& ctx) {
         TThis::PartitionHelper.Close(ctx);
-        if (NKikimrPQ::ETopicPartitionStatus::Active == ev->Get()->Record.GetStatus() 
+        if (NKikimrPQ::ETopicPartitionStatus::Active == ev->Get()->Record.GetStatus()
                 && ev->Get()->Record.HasSeqNo()
                 && ev->Get()->Record.GetSeqNo() > 0) {
             // Fast path: the partition ative and already written