From 8d4b5186eab6b0f9ad9a26e8ae01a4ce54641a9c Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Tue, 28 Apr 2026 19:52:59 +0800 Subject: [PATCH 1/7] lightning: add cloud worker write timeout handling --- pkg/lightning/backend/local/job_worker.go | 21 ++++++++++++++++++- .../backend/local/job_worker_test.go | 16 ++++++++++++++ 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/pkg/lightning/backend/local/job_worker.go b/pkg/lightning/backend/local/job_worker.go index 21dc25ccfef47..cdc31b1bec481 100644 --- a/pkg/lightning/backend/local/job_worker.go +++ b/pkg/lightning/backend/local/job_worker.go @@ -16,6 +16,7 @@ package local import ( "context" + goerrors "errors" "io" "strings" "sync" @@ -323,7 +324,25 @@ func (*objStoreRegionJobWorker) preRunJob(_ context.Context, _ *regionJob) error } // we don't need to limit write speed as we write to tikv-worker. -func (w *objStoreRegionJobWorker) write(ctx context.Context, job *regionJob) (*tikvWriteResult, error) { +func (w *objStoreRegionJobWorker) write(ctx context.Context, job *regionJob) (ret *tikvWriteResult, err error) { + var cancel context.CancelFunc + timeout := 15 * time.Minute + ctx, cancel = context.WithTimeoutCause(ctx, timeout, common.ErrWriteTooSlow) + defer cancel() + + 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-worker") + err = errors.Trace(cause) + } + } + }() + firstKey, _, err := job.ingestData.GetFirstAndLastKey(job.keyRange.Start, job.keyRange.End) if err != nil { return nil, errors.Trace(err) diff --git a/pkg/lightning/backend/local/job_worker_test.go b/pkg/lightning/backend/local/job_worker_test.go index bd2acbf884cb7..8d85024a5a2a0 100644 --- a/pkg/lightning/backend/local/job_worker_test.go +++ b/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,21 @@ func TestCloudRegionJobWorker(t *testing.T) { require.True(t, ctrl.Satisfied()) }) + t.Run("timeout while creating ingest client should return retryable error", func(t *testing.T) { + job := ®ionJob{ + keyRange: engineapi.Range{Start: []byte("a"), End: []byte("z")}, + stage: regionScanned, + ingestData: mockIngestData{{[]byte("a"), []byte("a")}}, + } + mockIngestCli.EXPECT().WriteClient(gomock.Any(), gomock.Any()).Return(nil, context.DeadlineExceeded) + ctx, cancel := context.WithTimeoutCause(context.Background(), 0, common.ErrWriteTooSlow) + defer cancel() + writeRes, err := cloudW.write(ctx, job) + require.ErrorIs(t, err, common.ErrWriteTooSlow) + require.Nil(t, writeRes) + require.True(t, ctrl.Satisfied()) + }) + t.Run("failed to write data", func(t *testing.T) { job := ®ionJob{ keyRange: engineapi.Range{Start: []byte("a"), End: []byte("z")}, From f2f973981fce76fe4d39a54c339404cd66cace84 Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Tue, 28 Apr 2026 20:01:59 +0800 Subject: [PATCH 2/7] lightning: unify region job write timeout handling --- pkg/lightning/backend/local/job_worker.go | 43 ++++++++++--------- .../backend/local/job_worker_test.go | 12 ++++-- pkg/lightning/backend/local/region_job.go | 29 +------------ 3 files changed, 32 insertions(+), 52 deletions(-) diff --git a/pkg/lightning/backend/local/job_worker.go b/pkg/lightning/backend/local/job_worker.go index cdc31b1bec481..39ff73a275b48 100644 --- a/pkg/lightning/backend/local/job_worker.go +++ b/pkg/lightning/backend/local/job_worker.go @@ -198,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) { @@ -267,6 +267,27 @@ func (w *regionJobBaseWorker) runJob(ctx context.Context, job *regionJob) error } } +func (w *regionJobBaseWorker) writeWithTimeout( + ctx context.Context, + job *regionJob, +) (ret *tikvWriteResult, err error) { + timeout := 15 * time.Minute + 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 @@ -324,25 +345,7 @@ func (*objStoreRegionJobWorker) preRunJob(_ context.Context, _ *regionJob) error } // we don't need to limit write speed as we write to tikv-worker. -func (w *objStoreRegionJobWorker) write(ctx context.Context, job *regionJob) (ret *tikvWriteResult, err error) { - var cancel context.CancelFunc - timeout := 15 * time.Minute - ctx, cancel = context.WithTimeoutCause(ctx, timeout, common.ErrWriteTooSlow) - defer cancel() - - 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-worker") - err = errors.Trace(cause) - } - } - }() - +func (w *objStoreRegionJobWorker) write(ctx context.Context, job *regionJob) (*tikvWriteResult, error) { firstKey, _, err := job.ingestData.GetFirstAndLastKey(job.keyRange.Start, job.keyRange.End) if err != nil { return nil, errors.Trace(err) diff --git a/pkg/lightning/backend/local/job_worker_test.go b/pkg/lightning/backend/local/job_worker_test.go index 8d85024a5a2a0..43851dc26ea5b 100644 --- a/pkg/lightning/backend/local/job_worker_test.go +++ b/pkg/lightning/backend/local/job_worker_test.go @@ -306,18 +306,22 @@ func TestCloudRegionJobWorker(t *testing.T) { require.True(t, ctrl.Satisfied()) }) - t.Run("timeout while creating ingest client should return retryable error", func(t *testing.T) { + t.Run("timeout while creating ingest client should be handled by runJob wrapper", func(t *testing.T) { job := ®ionJob{ 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() - writeRes, err := cloudW.write(ctx, job) - require.ErrorIs(t, err, common.ErrWriteTooSlow) - require.Nil(t, writeRes) + 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()) }) diff --git a/pkg/lightning/backend/local/region_job.go b/pkg/lightning/backend/local/region_job.go index bd610786078fc..9e649cbfb540e 100644 --- a/pkg/lightning/backend/local/region_job.go +++ b/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 @@ -486,6 +458,7 @@ func (local *Backend) doWrite(ctx context.Context, j *regionJob) (ret *tikvWrite } // preparation work for the write timeout fault injection, only enabled if the following failpoint is enabled + wctx := ctx wcancel := func() {} failpoint.Inject("shortWaitNTimeout", func(val failpoint.Value) { var innerTimeout time.Duration From 4a06c1f71e122fb74aa01800c1b7dbcc931c19a8 Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Wed, 6 May 2026 12:03:32 +0800 Subject: [PATCH 3/7] change --- .../tests/lightning_write_timeout/run.sh | 2 +- pkg/lightning/backend/local/job_worker.go | 10 +++++++- pkg/lightning/backend/local/region_job.go | 25 +++---------------- 3 files changed, 13 insertions(+), 24 deletions(-) diff --git a/lightning/tests/lightning_write_timeout/run.sh b/lightning/tests/lightning_write_timeout/run.sh index e291b2b635e48..ce305be4cd461 100644 --- a/lightning/tests/lightning_write_timeout/run.sh +++ b/lightning/tests/lightning_write_timeout/run.sh @@ -45,4 +45,4 @@ 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' +check_lightning_log_contains 'Experiencing a wait timeout while writing to TiKV' diff --git a/pkg/lightning/backend/local/job_worker.go b/pkg/lightning/backend/local/job_worker.go index 39ff73a275b48..1e7267ebf127c 100644 --- a/pkg/lightning/backend/local/job_worker.go +++ b/pkg/lightning/backend/local/job_worker.go @@ -275,13 +275,21 @@ func (w *regionJobBaseWorker) writeWithTimeout( ctx, cancel := context.WithTimeoutCause(ctx, timeout, common.ErrWriteTooSlow) defer cancel() + wcancel := func() {} + failpoint.Inject("shortWaitNTimeout", func(val failpoint.Value) { + innerTimeout := time.Duration(val.(int)) * time.Millisecond + tidblogutil.Logger(ctx).Info("Injecting a timeout to write context.") + ctx, wcancel = context.WithTimeoutCause(ctx, innerTimeout, common.ErrWriteTooSlow) + }) + defer wcancel() + 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") + tidblogutil.Logger(ctx).Info("Experiencing a wait timeout while writing to TiKV") err = errors.Trace(cause) } } diff --git a/pkg/lightning/backend/local/region_job.go b/pkg/lightning/backend/local/region_job.go index 9e649cbfb540e..80696419c8453 100644 --- a/pkg/lightning/backend/local/region_job.go +++ b/pkg/lightning/backend/local/region_job.go @@ -457,20 +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 - wctx := ctx - 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{} @@ -483,7 +469,7 @@ func (local *Backend) doWrite(ctx context.Context, j *regionJob) (ret *tikvWrite 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 @@ -493,13 +479,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 { From b5434fd5ee69d643e1d00038479971790dac6e46 Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Wed, 6 May 2026 13:32:41 +0800 Subject: [PATCH 4/7] lint --- pkg/lightning/backend/local/job_worker.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pkg/lightning/backend/local/job_worker.go b/pkg/lightning/backend/local/job_worker.go index 1e7267ebf127c..38f08e157e705 100644 --- a/pkg/lightning/backend/local/job_worker.go +++ b/pkg/lightning/backend/local/job_worker.go @@ -277,7 +277,8 @@ func (w *regionJobBaseWorker) writeWithTimeout( wcancel := func() {} failpoint.Inject("shortWaitNTimeout", func(val failpoint.Value) { - innerTimeout := time.Duration(val.(int)) * time.Millisecond + ms := val.(int) + innerTimeout := time.Duration(ms) * time.Millisecond tidblogutil.Logger(ctx).Info("Injecting a timeout to write context.") ctx, wcancel = context.WithTimeoutCause(ctx, innerTimeout, common.ErrWriteTooSlow) }) From 8f2b2ef8804a560d1452b38255174ecf04a6422c Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Wed, 6 May 2026 13:34:59 +0800 Subject: [PATCH 5/7] comment --- pkg/lightning/backend/local/job_worker.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pkg/lightning/backend/local/job_worker.go b/pkg/lightning/backend/local/job_worker.go index 38f08e157e705..765e2d35e9e8f 100644 --- a/pkg/lightning/backend/local/job_worker.go +++ b/pkg/lightning/backend/local/job_worker.go @@ -271,6 +271,9 @@ 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 ctx, cancel := context.WithTimeoutCause(ctx, timeout, common.ErrWriteTooSlow) defer cancel() From 122678527160a59e2c639adb298083778512cce0 Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Wed, 6 May 2026 16:13:20 +0800 Subject: [PATCH 6/7] fix comments --- lightning/tests/lightning_write_timeout/run.sh | 2 +- pkg/lightning/backend/local/job_worker.go | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/lightning/tests/lightning_write_timeout/run.sh b/lightning/tests/lightning_write_timeout/run.sh index ce305be4cd461..881a3f1e17dcc 100644 --- a/lightning/tests/lightning_write_timeout/run.sh +++ b/lightning/tests/lightning_write_timeout/run.sh @@ -45,4 +45,4 @@ 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' +check_lightning_log_contains 'experiencing a wait timeout while writing to TiKV' diff --git a/pkg/lightning/backend/local/job_worker.go b/pkg/lightning/backend/local/job_worker.go index 765e2d35e9e8f..246a12d701869 100644 --- a/pkg/lightning/backend/local/job_worker.go +++ b/pkg/lightning/backend/local/job_worker.go @@ -282,7 +282,7 @@ func (w *regionJobBaseWorker) writeWithTimeout( failpoint.Inject("shortWaitNTimeout", func(val failpoint.Value) { ms := val.(int) innerTimeout := time.Duration(ms) * time.Millisecond - tidblogutil.Logger(ctx).Info("Injecting a timeout to write context.") + tidblogutil.Logger(ctx).Info("injecting a timeout to write context.") ctx, wcancel = context.WithTimeoutCause(ctx, innerTimeout, common.ErrWriteTooSlow) }) defer wcancel() @@ -293,7 +293,7 @@ func (w *regionJobBaseWorker) writeWithTimeout( } 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") + tidblogutil.Logger(ctx).Info("experiencing a wait timeout while writing to TiKV") err = errors.Trace(cause) } } From f4cb5fbeb3543fea86f00993f277633f206ad867 Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Wed, 6 May 2026 17:23:12 +0800 Subject: [PATCH 7/7] fix test --- lightning/tests/lightning_write_timeout/run.sh | 2 +- pkg/lightning/backend/local/job_worker.go | 11 +++-------- pkg/lightning/backend/local/region_job.go | 6 +++++- 3 files changed, 9 insertions(+), 10 deletions(-) diff --git a/lightning/tests/lightning_write_timeout/run.sh b/lightning/tests/lightning_write_timeout/run.sh index 881a3f1e17dcc..56bd3d728e0a3 100644 --- a/lightning/tests/lightning_write_timeout/run.sh +++ b/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' diff --git a/pkg/lightning/backend/local/job_worker.go b/pkg/lightning/backend/local/job_worker.go index 246a12d701869..1b25153a89317 100644 --- a/pkg/lightning/backend/local/job_worker.go +++ b/pkg/lightning/backend/local/job_worker.go @@ -275,17 +275,12 @@ func (w *regionJobBaseWorker) writeWithTimeout( // 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() - - wcancel := func() {} failpoint.Inject("shortWaitNTimeout", func(val failpoint.Value) { ms := val.(int) - innerTimeout := time.Duration(ms) * time.Millisecond - tidblogutil.Logger(ctx).Info("injecting a timeout to write context.") - ctx, wcancel = context.WithTimeoutCause(ctx, innerTimeout, common.ErrWriteTooSlow) + timeout = time.Duration(ms) * time.Millisecond }) - defer wcancel() + ctx, cancel := context.WithTimeoutCause(ctx, timeout, common.ErrWriteTooSlow) + defer cancel() ret, err = w.writeFn(ctx, job) if err == nil { diff --git a/pkg/lightning/backend/local/region_job.go b/pkg/lightning/backend/local/region_job.go index 80696419c8453..ac9a48888aa22 100644 --- a/pkg/lightning/backend/local/region_job.go +++ b/pkg/lightning/backend/local/region_job.go @@ -465,7 +465,11 @@ 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