fix: mongodb restore and pre-terminate hang (#2335)

gnolong · apecloud-bot · commit 98744355de1f · 2025-12-24T06:21:49.000Z
(cherry picked from commit 01f044f)
diff --git a/addons/mongodb/dataprotection/common-scripts.sh b/addons/mongodb/dataprotection/common-scripts.sh
@@ -366,3 +366,66 @@ function process_restore_end_signal() {
     done
     echo "INFO: Prepare restore end signal completed."
 }
+
+function get_describe_backup_info() {
+  max_retries=360
+  retry_interval=2
+  attempt=1
+  describe_result=""
+  set +e
+  while [ $attempt -le $max_retries ]; do
+      describe_result=$(pbm describe-backup --mongodb-uri "$PBM_MONGODB_URI" "$backup_name" -o json 2>&1)
+      if [ $? -eq 0 ] && [ -n "$describe_result" ]; then
+          break
+      elif echo "$describe_result" | grep -q "not found"; then
+          echo "INFO: Attempt $attempt: backup $backup_name not found, retrying in ${retry_interval}s..."
+          if [ $((attempt % 30)) -eq 29 ]; then
+              echo "INFO: Sync PBM config from storage again."
+              sync_pbm_config_from_storage
+          fi
+          sleep $retry_interval
+          ((attempt++))
+          continue
+      else
+          echo "ERROR: Failed to get backup metadata: $describe_result"
+          exit 1
+      fi
+  done
+  set -e
+
+  if [ -z "$describe_result" ] || echo "$describe_result" | grep -q "not found"; then
+      echo "ERROR: Failed to get backup metadata after $max_retries attempts"
+      exit 1
+  fi
+}
+
+function wait_for_restoring() {
+  local cnf_file="${MOUNT_DIR}/tmp/pbm_restore.cnf"
+  cat <<EOF > ${MOUNT_DIR}/tmp/pbm_restore.cnf
+storage:
+  type: s3
+  s3:
+    region: ${S3_REGION}
+    bucket: ${S3_BUCKET}
+    prefix: ${S3_PREFIX}
+    endpointUrl: ${S3_ENDPOINT}
+    forcePathStyle: ${S3_FORCE_PATH_STYLE:-false}
+    credentials:
+      access-key-id: ${S3_ACCESS_KEY}
+      secret-access-key: ${S3_SECRET_KEY}
+EOF
+
+  while true; do
+    restore_status=$(pbm describe-restore "$restore_name" -c $cnf_file -o json | jq -r '.status') 
+    echo "INFO: Restore $restore_name status: $restore_status"
+    if [ "$restore_status" = "done" ]; then
+      rm $cnf_file
+      break
+    elif [ "$restore_status" = "" ] || [ "$restore_status" = "starting" ] || [ "$restore_status" = "running" ]; then
+      sleep 5
+    else
+      rm $cnf_file
+      exit 1
+    fi
+  done
+}
diff --git a/addons/mongodb/dataprotection/rs-pbm-full-restore.sh b/addons/mongodb/dataprotection/rs-pbm-full-restore.sh
@@ -37,31 +37,7 @@ if [ -z "$backup_type" ] || [ -z "$backup_name" ]; then
     exit 1
 fi
 
-MAX_RETRIES=360
-RETRY_INTERVAL=2
-attempt=1
-describe_result=""
-set +e
-while [ $attempt -le $MAX_RETRIES ]; do
-    describe_result=$(pbm describe-backup --mongodb-uri "$PBM_MONGODB_URI" "$backup_name" -o json 2>&1)
-    if [ $? -eq 0 ] && [ -n "$describe_result" ]; then
-        break
-    elif echo "$describe_result" | grep -q "not found"; then
-        echo "INFO: Attempt $attempt: Failed to get backup metadata, retrying in ${RETRY_INTERVAL}s..."
-        sleep $RETRY_INTERVAL
-        ((attempt++))
-        continue
-    else
-        echo "ERROR: Failed to get backup metadata: $describe_result"
-        exit 1
-    fi
-done
-set -e
-
-if [ -z "$describe_result" ] || echo "$describe_result" | grep -q "not found"; then
-    echo "ERROR: Failed to get backup metadata after $MAX_RETRIES attempts"
-    exit 1
-fi
+get_describe_backup_info
 
 rs_name=$(echo "$describe_result" | jq -r '.replsets[0].name')
 mappings="$MONGODB_REPLICA_SET_NAME=$rs_name"
@@ -71,6 +47,8 @@ process_restore_start_signal
 
 wait_for_other_operations
 
-pbm restore $backup_name --mongodb-uri "$PBM_MONGODB_URI" --replset-remapping "$mappings" --wait
+restore_name=$(pbm restore $backup_name --mongodb-uri "$PBM_MONGODB_URI" --replset-remapping "$mappings" -o json | jq -r '.name')
+
+wait_for_restoring
 
 process_restore_end_signal
diff --git a/addons/mongodb/dataprotection/rs-pbm-pitr-restore.sh b/addons/mongodb/dataprotection/rs-pbm-pitr-restore.sh
@@ -28,7 +28,6 @@ sync_pbm_storage_config
 
 sync_pbm_config_from_storage
 
-
 process_restore_start_signal
 
 extras=$(cat /dp_downward/status_extras)
@@ -43,7 +42,9 @@ echo "INFO: Starting restore..."
 
 wait_for_other_operations
 
-pbm restore --time="$recovery_target_time" --mongodb-uri "$PBM_MONGODB_URI" --replset-remapping "$mappings" --wait
+restore_name=$(pbm restore --time="$recovery_target_time" --mongodb-uri "$PBM_MONGODB_URI" --replset-remapping "$mappings" -o json | jq -r '.name')
+
+wait_for_restoring
 
 process_restore_end_signal
 
diff --git a/addons/mongodb/dataprotection/shard-pbm-full-restore.sh b/addons/mongodb/dataprotection/shard-pbm-full-restore.sh
@@ -37,31 +37,7 @@ if [ -z "$backup_type" ] || [ -z "$backup_name" ]; then
     exit 1
 fi
 
-MAX_RETRIES=360
-RETRY_INTERVAL=2
-attempt=1
-describe_result=""
-set +e
-while [ $attempt -le $MAX_RETRIES ]; do
-    describe_result=$(pbm describe-backup --mongodb-uri "$PBM_MONGODB_URI" "$backup_name" -o json 2>&1)
-    if [ $? -eq 0 ] && [ -n "$describe_result" ]; then
-        break
-    elif echo "$describe_result" | grep -q "not found"; then
-        echo "INFO: Attempt $attempt: Failed to get backup metadata, retrying in ${RETRY_INTERVAL}s..."
-        sleep $RETRY_INTERVAL
-        ((attempt++))
-        continue
-    else
-        echo "ERROR: Failed to get backup metadata: $describe_result"
-        exit 1
-    fi
-done
-set -e
-
-if [ -z "$describe_result" ] || echo "$describe_result" | grep -q "not found"; then
-    echo "ERROR: Failed to get backup metadata after $MAX_RETRIES attempts"
-    exit 1
-fi
+get_describe_backup_info
 
 configsvr_name=$(echo "$describe_result" | jq -r '.replsets[] | select(.configsvr == true) | .name')
 echo "INFO: Config server replica set name: $configsvr_name"
@@ -104,6 +80,8 @@ process_restore_start_signal
 
 wait_for_other_operations
 
-pbm restore $backup_name --mongodb-uri "$PBM_MONGODB_URI" --replset-remapping "$mappings" --wait
+restore_name=$(pbm restore $backup_name --mongodb-uri "$PBM_MONGODB_URI" --replset-remapping "$mappings" -o json | jq -r '.name')
+
+wait_for_restoring
 
 process_restore_end_signal
diff --git a/addons/mongodb/dataprotection/shard-pbm-pitr-restore.sh b/addons/mongodb/dataprotection/shard-pbm-pitr-restore.sh
@@ -76,7 +76,9 @@ echo "INFO: Starting restore..."
 
 wait_for_other_operations
 
-pbm restore --time="$recovery_target_time" --mongodb-uri "$PBM_MONGODB_URI" --replset-remapping "$mappings" --wait
+restore_name=$(pbm restore --time="$recovery_target_time" --mongodb-uri "$PBM_MONGODB_URI" --replset-remapping "$mappings" -o json | jq -r '.name')
+
+wait_for_restoring
 
 process_restore_end_signal
 
diff --git a/addons/mongodb/scripts/mongodb-shard-manage.sh b/addons/mongodb/scripts/mongodb-shard-manage.sh
@@ -31,7 +31,12 @@ wait_for_mongos() {
 check_shard_exists() {
     # check if the shard exists in the config database
     local shard_exists
-    shard_exists=$($CLUSTER_MONGO "db.getSiblingDB(\"config\").shards.find({ _id: \"$MONGODB_REPLICA_SET_NAME\" })")
+    shard_exists=$($CLUSTER_MONGO "db.getSiblingDB(\"config\").shards.find({ _id: \"$MONGODB_REPLICA_SET_NAME\" })" 2>/dev/null)
+    if [ $? -ne 0 ]; then
+        echo "ERROR: Failed to check if shard $MONGODB_REPLICA_SET_NAME exists." >&2
+        exit 1
+    fi
+    echo "INFO: Check if shard $MONGODB_REPLICA_SET_NAME exists: $shard_exists"
     if [ -n "$shard_exists" ]; then
         return 0 # true
     else
@@ -63,7 +68,11 @@ initialize_or_scale_out_mongodb_shard() {
 get_remove_shard_status() {
     # Execute the removeShard command and capture its JSON output
     local result
-    result=$($CLUSTER_MONGO "EJSON.stringify(db.adminCommand( { removeShard: \"$MONGODB_REPLICA_SET_NAME\" } ))")
+    if [ "$CLIENT" = "mongosh" ]; then
+        result=$($CLUSTER_MONGO "EJSON.stringify(db.adminCommand( { removeShard: \"$MONGODB_REPLICA_SET_NAME\" } ))")
+    else
+        result=$($CLUSTER_MONGO "JSON.stringify(db.adminCommand( { removeShard: \"$MONGODB_REPLICA_SET_NAME\" } ))")
+    fi
     echo "$result"
 }
 
@@ -76,6 +85,32 @@ get_remove_shard_state() {
     echo "$state"
 }
 
+get_remaining_jumbo_chunks() {
+    local result=$1
+    # Parse and log the jumboChunks count using jq
+    local jumbo_chunks
+    if [ "$CLIENT" = "mongosh" ]; then
+        jumbo_chunks=$(echo "$result" | jq -r '.remaining.jumboChunks // 0')
+    else
+        jumbo_chunks=$(echo "$result" | jq -r '.remaining.jumboChunks.numberLong // 0')
+    fi
+    # Return the jumboChunks count as the function output
+    echo "$jumbo_chunks"
+}
+
+get_remaining_chunks() {
+    local result=$1
+    # Parse and log the chunks count using jq
+    local chunks
+    if [ "$CLIENT" = "mongosh" ]; then
+        chunks=$(echo "$result" | jq -r '.remaining.chunks // 0')
+    else
+        chunks=$(echo "$result" | jq -r '.remaining.chunks.numberLong // 0')
+    fi
+    # Return the chunks count as the function output
+    echo "$chunks"
+}
+
 delete_or_scale_in_mongodb_shard() {
     # Check if the shard is scaling in
     if [[ $KB_CLUSTER_COMPONENT_IS_SCALING_IN != "true" ]]; then
@@ -111,13 +146,13 @@ delete_or_scale_in_mongodb_shard() {
         if [ "$state" = "completed" ]; then
             break
         elif [ "$state" = "ongoing" ]; then
-            remaining_jumboChunks=$(echo "$status_json" | jq -r '.remaining.jumboChunks')
+            remaining_jumboChunks=$(get_remaining_jumbo_chunks "$status_json")
             if [ "$remaining_jumboChunks" -gt 0 ]; then
                 echo "INFO: $remaining_jumboChunks jumbo chunks remaining, please clear jumbo chunks before removing the shard."
                 exit 1
             fi
 
-            remaining_chunks=$(echo "$status_json" | jq -r '.remaining.chunks')
+            remaining_chunks=$(get_remaining_chunks "$status_json")
             echo "INFO: $remaining_chunks chunks remaining."
             if [ "$remaining_chunks" -eq 0 ]; then
                 dbs_to_move=$(echo "$status_json" | jq -r '.dbsToMove[]')
