This is an automated cherry-pick of #61346

Signed-off-by: ti-chi-bot <[email protected]>

Author: OliverS929

br/pkg/lightning/backend/local/region_job.go

@@ -204,11 +204,15 @@ func (local *Backend) writeToTiKV(ctx context.Context, j *regionJob) error {
 	return err
 }
 
+<<<<<<< HEAD:br/pkg/lightning/backend/local/region_job.go
 func (local *Backend) doWrite(ctx context.Context, j *regionJob) error {
 	if j.stage != regionScanned {
 		return nil
 	}
 
+=======
+func (local *Backend) doWrite(ctx context.Context, j *regionJob) (ret *tikvWriteResult, err error) {
+>>>>>>> cc0925eeafc (Lightning: Attempt to return writeTooSlow when we experience write timeout. (#61346)):pkg/lightning/backend/local/region_job.go
 	failpoint.Inject("fakeRegionJobs", func() {
 		front := j.injected[0]
 		j.injected = j.injected[1:]
@@ -221,9 +225,33 @@ func (local *Backend) doWrite(ctx context.Context, j *regionJob) error {
 	})
 
 	var cancel context.CancelFunc
-	ctx, cancel = context.WithTimeoutCause(ctx, 15*time.Minute, common.ErrWriteTooSlow)
+	// set a timeout for the write operation, if it takes too long, we will return with common.ErrWriteTooSlow and let caller retry the whole job instead of being stuck forever.
+	timeout := 15 * time.Minute
+	ctx, cancel = context.WithTimeoutCause(ctx, timeout, common.ErrWriteTooSlow)
 	defer cancel()
 
+	// A defer function to handle all DeadlineExceeded errors that may occur
+	// during the write operation using this context with 15 minutes timeout.
+	// When the error is "context deadline exceeded", we will check if the cause
+	// is common.ErrWriteTooSlow and return the common.ErrWriteTooSlow instead so
+	// our caller would be able to retry this doWrite operation. By doing this
+	// defer we are hoping to handle all DeadlineExceeded error during this
+	// write, either from gRPC stream or write limiter WaitN operation.
+	wctx := ctx
+	defer func() {
+		if err == nil {
+			return
+		}
+		if errors.Cause(err) == context.DeadlineExceeded {
+			if cause := context.Cause(wctx); goerrors.Is(cause, common.ErrWriteTooSlow) {
+				log.FromContext(ctx).Info("Experiencing a wait timeout while writing to tikv",
+					zap.Int("store-write-bwlimit", local.BackendConfig.StoreWriteBWLimit),
+					zap.Int("limit-size", local.writeLimiter.Limit()))
+				err = errors.Trace(cause) // return the common.ErrWriteTooSlow instead to let caller retry it
+			}
+		}
+	}()
+
 	apiVersion := local.tikvCodec.GetAPIVersion()
 	clientFactory := local.importClientFactory
 	kvBatchSize := local.KVWriteBatchSize
@@ -341,6 +369,19 @@ func (local *Backend) doWrite(ctx context.Context, j *regionJob) error {
 		regionMaxSize = j.regionSplitSize * 4 / 3
 	}
 
+	// preparation work for the write timeout fault injection, only enabled if the following failpoint is enabled
+	wcancel := func() {}
+	failpoint.Inject("shortWaitNTimeout", func(val failpoint.Value) {
+		var innerTimeout time.Duration
+		// GO_FAILPOINTS action supplies the duration in
+		ms, _ := val.(int)
+		innerTimeout = time.Duration(ms) * time.Millisecond
+		log.FromContext(ctx).Info("Injecting a timeout to write context.")
+		wctx, wcancel = context.WithTimeoutCause(
+			ctx, innerTimeout, common.ErrWriteTooSlow)
+	})
+	defer wcancel()
+
 	flushKVs := func() error {
 		req.Chunk.(*sst.WriteRequest_Batch).Batch.Pairs = pairs[:count]
 		preparedMsg := &grpc.PreparedMsg{}
@@ -351,7 +392,25 @@ func (local *Backend) doWrite(ctx context.Context, j *regionJob) error {
 		}
 
 		for i := range clients {
-			if err := writeLimiter.WaitN(ctx, allPeers[i].StoreId, int(size)); err != nil {
+			// original ctx would be used when failpoint is not enabled
+			// that new context would be used when failpoint is enabled
+			err := writeLimiter.WaitN(wctx, allPeers[i].StoreId, int(size))
+			if err != nil {
+				// We expect to encounter two types of errors here:
+				// 1. context.DeadlineExceeded — occurs when the calculated delay is
+				//    less than the remaining time in the context, but the context
+				//    expires while sleeping.
+				// 2. "rate: Wait(n=%d) would exceed context deadline" — a fast-fail
+				//    path triggered when the delay already exceeds the remaining
+				//    time for context before sleeping.
+				//
+				// Unfortunately, we cannot precisely control when the context will
+				// expire, so both scenarios are valid and expected.
+				// Fortunately, the "rate: Wait" error is already treated as
+				// retryable, so we only need to explicitly handle
+				// context.DeadlineExceeded here.
+				// We rely on the defer function at the top of doWrite to handle it
+				// for us in general.
 				return errors.Trace(err)
 			}
 			if err := clients[i].SendMsg(preparedMsg); err != nil {

br/tests/_utils/check_lightning_log_contains

@@ -0,0 +1,30 @@
+#!/bin/sh
+#
+# Copyright 2025 PingCAP, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -eu
+
+log_file="${TEST_DIR}/lightning.log"
+if [ $# -gt 1 ]; then
+    log_file=$2
+fi
+
+if ! grep -Fq "$1" "$log_file"; then
+    echo "TEST FAILED: LIGHTNING LOG DOES NOT CONTAIN '$1'"
+    echo "____________________________________"
+    cat "$log_file"
+    echo "^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^"
+    exit 1
+fi

lightning/tests/README.md

@@ -0,0 +1,76 @@
+# Integration tests
+
+This folder contains all tests which relies on external processes such as TiDB.
+
+## Preparations
+
+1. The following executables must be copied or linked into these locations:
+
+    * `bin/tidb-server`
+    * `bin/tikv-server`
+    * `bin/pd-server`
+    * `bin/tiflash`
+    * `bin/minio`
+    * `bin/mc`
+
+    The versions must be ≥2.1.0.
+
+    **Only some tests requires `minio`/`mc`** which can be downloaded from the official site, you can skip them if you don't need to run those cases.
+
+    You can use `tiup` to download binaries related to TiDB cluster, and then link them to the `bin` directory:
+    ```shell
+    cluster_version=v8.1.0 # change to the version you need
+    tiup install tidb:$cluster_version tikv:$cluster_version pd:$cluster_version tiflash:$cluster_version
+    ln -s ~/.tiup/components/tidb/$cluster_version/tidb-server bin/tidb-server
+    ln -s ~/.tiup/components/tikv/$cluster_version/tikv-server bin/tikv-server
+    ln -s ~/.tiup/components/pd/$cluster_version/pd-server bin/pd-server
+    ln -s ~/.tiup/components/tiflash/$cluster_version/tiflash/tiflash bin/tiflash
+    ```
+
+2. `make build_for_lightning_integration_test`
+   
+    `make server` to build the latest TiDB server if your test requires it.
+
+3. The following programs must be installed:
+
+    * `mysql` (the CLI client)
+    * `curl`
+    * `openssl`
+    * `wget`
+    * `lsof`
+
+4. The user executing the tests must have permission to create the folder
+    `/tmp/lightning_test`. All test artifacts will be written into this folder.
+
+## Running
+
+Run `make lightning_integration_test` to execute all the integration tests.
+- Logs will be written into `/tmp/lightning_test` directory.
+
+Run `tests/run.sh --debug` to pause immediately after all servers are started.
+
+If you only want to run some tests, you can use:
+```shell
+TEST_NAME="lightning_gcs lightning_view" lightning/tests/run.sh
+```
+
+Case names are separated by spaces.
+
+## Writing new tests
+
+1. New integration tests can be written as shell scripts in `tests/TEST_NAME/run.sh`.
+    - `TEST_NAME` should start with `lightning_`.
+    - The script should exit with a nonzero error code on failure.
+2. Add TEST_NAME to existing group in [run_group_lightning_tests.sh](./run_group_lightning_tests.sh)(Recommended), or add a new group for it.
+3. If you add a new group, the name of the new group must be added to CI  [lightning-integration-test](https://github.com/PingCAP-QE/ci/blob/main/pipelines/pingcap/tidb/latest/pull_lightning_integration_test.groovy).
+
+Several convenient commands are provided in [utils](../../tests/_utils/):
+
+* `run_sql <SQL>` — Executes an SQL query on the TiDB database
+* `run_lightning [CONFIG]` — Starts `tidb-lightning` using `tests/TEST_NAME/CONFIG.toml`
+* `check_contains <TEXT>` — Checks if the previous `run_sql` result contains the given text
+    (in `-E` format)
+* `check_not_contains <TEXT>` — Checks if the previous `run_sql` result does not contain the given
+    text (in `-E` format)
+* `check_lightning_log_contains <TEXT>` — Checks if the current lightning log contains the given text
+    (in `-E` format)

lightning/tests/lightning_write_timeout/config.toml

@@ -0,0 +1,5 @@
+[tikv-importer]
+store-write-bwlimit = "1MiB"
+
+[mydumper.csv]
+header = false

lightning/tests/lightning_write_timeout/run.sh

@@ -0,0 +1,48 @@
+#!/bin/bash
+#
+# Copyright 2025 PingCAP, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -eux
+
+CUR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+
+mkdir -p "$TEST_DIR/data"
+
+run_sql "DROP DATABASE IF EXISTS test;"
+run_sql "DROP TABLE IF EXISTS test.t;"
+
+cat <<EOF >"$TEST_DIR/data/test-schema-create.sql"
+CREATE DATABASE test;
+EOF
+cat <<EOF >"$TEST_DIR/data/test.t-schema.sql"
+CREATE TABLE test.t (
+  id int,
+  a int,
+  b int,
+  c int
+);
+EOF
+
+# Generate 200k rows. Total size is about 5MiB.
+set +x
+for i in {1..200000}; do
+  echo "$i,$i,$i,$i" >>"$TEST_DIR/data/test.t.0.csv"
+done
+set -x
+
+export GO_FAILPOINTS="github.com/pingcap/tidb/pkg/lightning/backend/local/shortWaitNTimeout=100*return(1)" 
+
+run_lightning --backend local -d "$TEST_DIR/data" --config "$CUR/config.toml"
+check_lightning_log_contains 'Experiencing a wait timeout while writing to tikv'

lightning/tests/run_group_lightning_tests.sh

@@ -0,0 +1,84 @@
+#!/usr/bin/env bash
+#
+# Copyright 2022 PingCAP, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# This script split the integration tests into 9 groups to support parallel group tests execution.
+# all the integration tests are located in br/tests directory. only the directories
+# containing run.sh will be considered as valid lightning integration tests. the script will print the total case number
+
+set -eo pipefail
+
+# Step 1
+CUR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+group=$1
+export COV_DIR="/tmp/group_cover"
+rm -rf $COV_DIR
+mkdir -p $COV_DIR
+
+# Define groups
+# Note: If new group is added, the group name must also be added to CI
+# * https://github.com/PingCAP-QE/ci/blob/main/pipelines/pingcap/tidb/latest/pull_lightning_integration_test.groovy
+# Each group of tests consumes as much time as possible, thus reducing CI waiting time.
+# Putting multiple light tests together and heavy tests in a separate group.
+declare -A groups
+groups=(
+	["G00"]='lightning_auto_random_default lightning_bom_file lightning_character_sets lightning_check_partial_imported lightning_checkpoint lightning_checkpoint_chunks lightning_checkpoint_columns lightning_checkpoint_dirty_tableid'
+	["G01"]='lightning_checkpoint_engines lightning_checkpoint_engines_order lightning_checkpoint_error_destroy lightning_checkpoint_parquet lightning_checkpoint_timestamp lightning_checksum_mismatch lightning_cmdline_override lightning_column_permutation lightning_common_handle lightning_compress lightning_concurrent-restore'
+	["G02"]='lightning_config_max_error lightning_config_skip_csv_header lightning_csv lightning_default-columns lightning_disable_scheduler_by_key_range lightning_disk_quota lightning_distributed_import lightning_drop_other_tables_halfway lightning_duplicate_detection lightning_duplicate_detection_new lightning_duplicate_resolution_error lightning_duplicate_resolution_error_pk_multiple_files lightning_duplicate_resolution_error_uk_multiple_files lightning_duplicate_resolution_error_uk_multiple_files_multicol_index lightning_duplicate_resolution_incremental'
+	["G03"]='lightning_duplicate_resolution_merge lightning_duplicate_resolution_replace_multiple_keys_clustered_pk lightning_duplicate_resolution_replace_multiple_keys_nonclustered_pk lightning_duplicate_resolution_replace_multiple_unique_keys_clustered_pk lightning_duplicate_resolution_replace_multiple_unique_keys_nonclustered_pk lightning_duplicate_resolution_replace_one_key lightning_duplicate_resolution_replace_one_key_multiple_conflicts_clustered_pk lightning_duplicate_resolution_replace_one_key_multiple_conflicts_nonclustered_pk'
+	["G04"]='lightning_duplicate_resolution_replace_one_unique_key_clustered_pk lightning_duplicate_resolution_replace_one_unique_key_multiple_conflicts_clustered_pk lightning_duplicate_resolution_replace_one_unique_key_multiple_conflicts_nonclustered_pk lightning_duplicate_resolution_replace_one_unique_key_nonclustered_varchar_pk lightning_error_summary lightning_examples lightning_exotic_filenames lightning_extend_routes'
+	["G05"]='lightning_fail_fast lightning_fail_fast_on_nonretry_err lightning_file_routing lightning_foreign_key lightning_gcs lightning_generated_columns lightning_ignore_columns lightning_import_compress lightning_incremental lightning_issue_282 lightning_issue_40657 lightning_issue_410 lightning_issue_519 lightning_local_backend lightning_max_incr'
+	["G06"]='lightning_max_random lightning_multi_valued_index lightning_new_collation lightning_no_schema lightning_parquet lightning_partition_incremental lightning_partitioned-table lightning_record_network lightning_reload_cert lightning_restore lightning_routes lightning_routes_panic lightning_row-format-v2 lightning_s3'
+	["G07"]='lightning_shard_rowid lightning_source_linkfile lightning_sqlmode lightning_tidb_duplicate_data lightning_tidb_rowid lightning_tiflash lightning_tikv_multi_rocksdb lightning_too_many_columns lightning_tool_135'
+	["G08"]='lightning_pd_leader_switch lightning_tool_1420 lightning_tool_1472 lightning_tool_241 lightning_ttl lightning_unused_config_keys lightning_various_types lightning_view lightning_write_batch lightning_write_limit lightning_write_timeout lightning_add_index lightning_alter_random lightning_auto_columns'
+)
+
+# Get other lightning cases not in groups, to avoid missing any case
+others=()
+for script in "$CUR"/*/run.sh; do
+	test_name="$(basename "$(dirname "$script")")"
+	if [[ $test_name != lightning_* ]]; then
+		continue
+	fi
+	# shellcheck disable=SC2076
+	if [[ ! " ${groups[*]} " =~ " ${test_name} " ]]; then
+		others=("${others[@]} ${test_name}")
+	fi
+done
+
+if [[ "$group" == "others" ]]; then
+	if [[ -z $others ]]; then
+		echo "All lightning test cases have been added to groups"
+		exit 0
+	fi
+	echo "Error: "$others" is not added to any group in br/tests/run_group_lightning_tests.sh"
+	exit 1
+elif [[ " ${!groups[*]} " =~ " ${group} " ]]; then
+	test_names="${groups[${group}]}"
+	# Run test cases
+	if [[ -n $test_names ]]; then
+		echo ""
+		echo "Run cases: ${test_names}"
+        for case_name in $test_names; do
+            echo "Run cases: ${case_name}"
+            rm -rf /tmp/lightning_test
+            mkdir -p /tmp/lightning_test
+            TEST_NAME=${case_name} ${CUR}/run.sh
+        done
+	fi
+else
+	echo "Error: invalid group name: ${group}"
+	exit 1
+fi