KIKIMR-13960: changes for out of space scenarios

4 files changed, 33 insertions, 4 deletions

diff --git a/ydb/core/kqp/executer_actor/kqp_data_executer.cpp b/ydb/core/kqp/executer_actor/kqp_data_executer.cpp
index fd6cbc846e7..383a2a34163 100644
--- a/ydb/core/kqp/executer_actor/kqp_data_executer.cpp
+++ b/ydb/core/kqp/executer_actor/kqp_data_executer.cpp
@@ -708,7 +708,12 @@ private:
                         case NKikimrTxDataShard::TError::SCHEME_ERROR:
                             return ReplyErrorAndDie(Ydb::StatusIds::SCHEME_ERROR, YqlIssue({},
                                 TIssuesIds::KIKIMR_SCHEME_MISMATCH, er.GetReason()));
-
+                        case NKikimrTxDataShard::TError::OUT_OF_SPACE:
+                        case NKikimrTxDataShard::TError::DISK_SPACE_EXHAUSTED: {
+                            auto issue = YqlIssue({}, TIssuesIds::KIKIMR_TEMPORARILY_UNAVAILABLE);
+                            AddDataShardErrors(result, issue);
+                            return ReplyErrorAndDie(Ydb::StatusIds::UNAVAILABLE, issue);
+                        }
                         default:
                             break;
                     }
diff --git a/ydb/core/kqp/ut/query/kqp_limits_ut.cpp b/ydb/core/kqp/ut/query/kqp_limits_ut.cpp
index 2fb47478608..e0e7d83aaba 100644
--- a/ydb/core/kqp/ut/query/kqp_limits_ut.cpp
+++ b/ydb/core/kqp/ut/query/kqp_limits_ut.cpp
@@ -13,6 +13,12 @@ using namespace NYdb::NTable;
 static const ui32 LargeTableShards = 8;
 static const ui32 LargeTableKeysPerShard = 1000000;
 
+namespace {
+    bool IsRetryable(const EStatus& status) {
+        return status == EStatus::OVERLOADED;    
+    }
+}
+
 static void CreateLargeTable(TKikimrRunner& kikimr, ui32 rowsPerShard, ui32 keyTextSize,
     ui32 dataTextSize, ui32 batchSizeRows = 100, ui32 fillShardsCount = LargeTableShards)
 {
@@ -193,7 +199,7 @@ Y_UNIT_TEST_SUITE(KqpLimits) {
             rowsBuilder.EndList();
 
             auto result = client.BulkUpsert("/Root/LargeTable", rowsBuilder.Build()).ExtractValueSync();
-            if (result.GetStatus() == EStatus::OVERLOADED) {
+            if (IsRetryable(result.GetStatus())) {
                 continue;
             } 
             if (result.GetStatus() != EStatus::SUCCESS) {
@@ -260,12 +266,13 @@ Y_UNIT_TEST_SUITE(KqpLimits) {
                 UPSERT INTO `/Root/LargeTable`
                 SELECT * FROM AS_TABLE($rows);
             )"), TTxControl::BeginTx().CommitTx(), paramsBuilder.Build()).ExtractValueSync();
-            if (result.GetStatus() == EStatus::OVERLOADED) {
+            if (IsRetryable(result.GetStatus())) {
                 continue;
             } 
             if (result.GetStatus() != EStatus::SUCCESS) {
                 result.GetIssues().PrintTo(Cerr);
                 UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::UNAVAILABLE, result.GetIssues().ToString());
+                UNIT_ASSERT_STRING_CONTAINS(result.GetIssues().ToString(), "OUT_OF_SPACE");
                 failedToInsert = true;
                 break;
             }
diff --git a/ydb/core/tx/datashard/datashard.cpp b/ydb/core/tx/datashard/datashard.cpp
index 4a2b8394a58..7f105b0bf85 100644
--- a/ydb/core/tx/datashard/datashard.cpp
+++ b/ydb/core/tx/datashard/datashard.cpp
@@ -2349,6 +2349,7 @@ Ydb::StatusIds::StatusCode ConvertToYdbStatusCode(NKikimrTxDataShard::TError::EK
         case NKikimrTxDataShard::TError::UNKNOWN:
         case NKikimrTxDataShard::TError::REPLY_SIZE_EXCEEDED:
         case NKikimrTxDataShard::TError::EXECUTION_CANCELLED:
+        case NKikimrTxDataShard::TError::DISK_SPACE_EXHAUSTED:
             return Ydb::StatusIds::INTERNAL_ERROR;
         case NKikimrTxDataShard::TError::BAD_ARGUMENT:
         case NKikimrTxDataShard::TError::READONLY:
diff --git a/ydb/core/tx/datashard/datashard__op_rows.cpp b/ydb/core/tx/datashard/datashard__op_rows.cpp
index de8ac3599f4..d732023d2c2 100644
--- a/ydb/core/tx/datashard/datashard__op_rows.cpp
+++ b/ydb/core/tx/datashard/datashard__op_rows.cpp
@@ -116,6 +116,10 @@ static void OutOfSpace(NKikimrTxDataShard::TEvUploadRowsResponse& response) {
     response.SetStatus(NKikimrTxDataShard::TError::OUT_OF_SPACE);
 }
 
+static void DiskSpaceExhausted(NKikimrTxDataShard::TEvUploadRowsResponse& response) {
+    response.SetStatus(NKikimrTxDataShard::TError::DISK_SPACE_EXHAUSTED);
+}
+
 static void WrongShardState(NKikimrTxDataShard::TEvUploadRowsResponse& response) {
     response.SetStatus(NKikimrTxDataShard::TError::WRONG_SHARD_STATE);
 }
@@ -129,6 +133,11 @@ static void OutOfSpace(NKikimrTxDataShard::TEvEraseRowsResponse& response) {
     response.SetStatus(NKikimrTxDataShard::TEvEraseRowsResponse::WRONG_SHARD_STATE);
 }
 
+static void DiskSpaceExhausted(NKikimrTxDataShard::TEvEraseRowsResponse& response) {
+    // NOTE: this function is never called, because erase is allowed when out of space
+    response.SetStatus(NKikimrTxDataShard::TEvEraseRowsResponse::WRONG_SHARD_STATE);
+}
+
 static void WrongShardState(NKikimrTxDataShard::TEvEraseRowsResponse& response) {
     response.SetStatus(NKikimrTxDataShard::TEvEraseRowsResponse::WRONG_SHARD_STATE);
 }
@@ -161,6 +170,8 @@ static bool MaybeReject(TDataShard* self, TEvRequest& ev, const TActorContext& c
     TString rejectReason;
     bool reject = self->CheckDataTxReject(txDesc, ctx, rejectStatus, rejectReason);
     bool outOfSpace = false;
+    bool isDiskSpaceExhausted = false;
+
 
     if (!reject && isWrite) {
         if (self->IsAnyChannelYellowStop()) {
@@ -172,6 +183,7 @@ static bool MaybeReject(TDataShard* self, TEvRequest& ev, const TActorContext& c
             reject = true;
             outOfSpace = true;
             rejectReason = "Cannot perform writes: database is out of disk space";
+            isDiskSpaceExhausted = true;
             self->IncCounter(COUNTER_PREPARE_OUT_OF_SPACE);
         }
     }
@@ -181,7 +193,11 @@ static bool MaybeReject(TDataShard* self, TEvRequest& ev, const TActorContext& c
     }
 
     if (outOfSpace) {
-        Reject<TEvResponse, TEvRequest>(self, ev, txDesc, rejectReason, &OutOfSpace, ctx);
+        if (isDiskSpaceExhausted) {
+            Reject<TEvResponse, TEvRequest>(self, ev, txDesc, rejectReason, &DiskSpaceExhausted, ctx);
+        } else {
+            Reject<TEvResponse, TEvRequest>(self, ev, txDesc, rejectReason, &OutOfSpace, ctx);
+        }
     } else {
         Reject<TEvResponse, TEvRequest>(self, ev, txDesc, rejectReason, &WrongShardState, ctx);
     }