localbackend: add write timeout handling for object store region jobs

lightning/tests/lightning_write_timeout/run.sh

@@ -42,7 +42,7 @@ for i in {1..200000}; do
 done
 set -x
 
-export GO_FAILPOINTS="github.com/pingcap/tidb/pkg/lightning/backend/local/shortWaitNTimeout=100*return(1)" 
+export GO_FAILPOINTS="github.com/pingcap/tidb/pkg/lightning/backend/local/shortWaitNTimeout=5*return(1000)"
 
 run_lightning --backend local -d "$TEST_DIR/data" --config "$CUR/config.toml"
-check_lightning_log_contains 'Experiencing a wait timeout while writing to tikv'
+check_lightning_log_contains 'experiencing a wait timeout while writing to TiKV'

pkg/lightning/backend/local/job_worker.go

@@ -16,6 +16,7 @@ package local
 
 import (
 	"context"
+	goerrors "errors"
 	"io"
 	"strings"
 	"sync"
@@ -197,7 +198,7 @@ func (w *regionJobBaseWorker) runJob(ctx context.Context, job *regionJob) error
 				job.convertStageTo(needRescan)
 				return nil
 			}
-			res, err := w.writeFn(ctx, job)
+			res, err := w.writeWithTimeout(ctx, job)
 			err = injectfailpoint.DXFRandomErrorWithOnePercentWrapper(err)
 			if err != nil {
 				if !w.isRetryableImportTiKVError(err) {
@@ -266,6 +267,34 @@ func (w *regionJobBaseWorker) runJob(ctx context.Context, job *regionJob) error
 	}
 }
 
+func (w *regionJobBaseWorker) writeWithTimeout(
+	ctx context.Context,
+	job *regionJob,
+) (ret *tikvWriteResult, err error) {
+	// 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
+	failpoint.Inject("shortWaitNTimeout", func(val failpoint.Value) {
+		ms := val.(int)
+		timeout = time.Duration(ms) * time.Millisecond
+	})
+	ctx, cancel := context.WithTimeoutCause(ctx, timeout, common.ErrWriteTooSlow)
+	defer cancel()
+
+	ret, err = w.writeFn(ctx, job)
+	if err == nil {
+		return ret, nil
+	}
+	if errors.Cause(err) == context.DeadlineExceeded {
+		if cause := context.Cause(ctx); goerrors.Is(cause, common.ErrWriteTooSlow) {
+			tidblogutil.Logger(ctx).Info("experiencing a wait timeout while writing to TiKV")
+			err = errors.Trace(cause)
+		}
+	}
+	return ret, err
+}
+
 func (*regionJobBaseWorker) isRetryableImportTiKVError(err error) bool {
 	err = errors.Cause(err)
 	// io.EOF is not retryable in normal case

pkg/lightning/backend/local/job_worker_test.go

@@ -28,6 +28,7 @@ import (
 	"github.com/pingcap/tidb/pkg/ingestor/errdef"
 	"github.com/pingcap/tidb/pkg/ingestor/ingestcli"
 	ingestclimock "github.com/pingcap/tidb/pkg/ingestor/ingestcli/mock"
+	"github.com/pingcap/tidb/pkg/lightning/common"
 	"github.com/pingcap/tidb/pkg/resourcemanager/pool/workerpool"
 	rcmgrutil "github.com/pingcap/tidb/pkg/resourcemanager/util"
 	"github.com/pingcap/tidb/pkg/testkit/testfailpoint"
@@ -305,6 +306,25 @@ func TestCloudRegionJobWorker(t *testing.T) {
 		require.True(t, ctrl.Satisfied())
 	})
 
+	t.Run("timeout while creating ingest client should be handled by runJob wrapper", func(t *testing.T) {
+		job := &regionJob{
+			keyRange:   engineapi.Range{Start: []byte("a"), End: []byte("z")},
+			stage:      regionScanned,
+			ingestData: mockIngestData{{[]byte("a"), []byte("a")}},
+			region: &split.RegionInfo{Region: &metapb.Region{
+				Id: 1, Peers: []*metapb.Peer{{StoreId: 1}},
+			}, Leader: &metapb.Peer{StoreId: 1}},
+		}
+		mockIngestCli.EXPECT().WriteClient(gomock.Any(), gomock.Any()).Return(nil, context.DeadlineExceeded)
+		ctx, cancel := context.WithTimeoutCause(context.Background(), 0, common.ErrWriteTooSlow)
+		defer cancel()
+		err := cloudW.runJob(ctx, job)
+		require.NoError(t, err)
+		require.Equal(t, needRescan, job.stage)
+		require.ErrorIs(t, job.lastRetryableErr, common.ErrWriteTooSlow)
+		require.True(t, ctrl.Satisfied())
+	})
+
 	t.Run("failed to write data", func(t *testing.T) {
 		job := &regionJob{
 			keyRange:   engineapi.Range{Start: []byte("a"), End: []byte("z")},

pkg/lightning/backend/local/region_job.go

@@ -339,34 +339,6 @@ func (local *Backend) doWrite(ctx context.Context, j *regionJob) (ret *tikvWrite
 		failpoint.Return(front.write.result, err)
 	})
 
-	var cancel context.CancelFunc
-	// 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) {
-				tidblogutil.Logger(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
@@ -485,19 +457,6 @@ func (local *Backend) doWrite(ctx context.Context, j *regionJob) (ret *tikvWrite
 		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
-		tidblogutil.Logger(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{}
@@ -506,11 +465,15 @@ func (local *Backend) doWrite(ctx context.Context, j *regionJob) (ret *tikvWrite
 		if err := preparedMsg.Encode(clients[0], req); err != nil {
 			return err
 		}
-
+		failpoint.Inject("shortWaitNTimeout", func(val failpoint.Value) {
+			// add 20 to make sure the ctx is timed out.
+			ms := val.(int) + 20
+			time.Sleep(time.Duration(ms) * time.Millisecond)
+		})
 		for i := range clients {
 			// 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))
+			err := writeLimiter.WaitN(ctx, 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
@@ -520,13 +483,8 @@ func (local *Backend) doWrite(ctx context.Context, j *regionJob) (ret *tikvWrite
 				//    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.
+				// "rate: Wait" error is already treated as retryable,
+				// writeWithTimeout will handle the context.DeadlineExceeded.
 				return errors.Trace(err)
 			}
 			if err := clients[i].SendMsg(preparedMsg); err != nil {