Lightning: Memory Leak on Large Source Files

[dsdashun]

Bug Report

Please answer these questions before submitting your issue. Thanks!

1. Minimal reproduce step (Required)

1. Prepare a large CSV file as the data source
2. Create the Lightning task with strict-mode set to false
3. Observe the memory usage of Lightning

2. What did you expect to see? (Required)

The memory usage should be stable, even if the source data is large, because Lightning is streaming-processing the data.

3. What did you see instead (Required)

The memory of Lightning process is increasing continuously. I suspect there might be memory leaks.

4. What is your TiDB version? (Required)

latest version

[dsdashun]

/component lightning

[dsdashun]

/severity major

[dsdashun]

Captured the heap profile:

[dsdashun]

I've got the memory leak point. It's in here:

tidb/br/pkg/lightning/backend/kv/session.go

Line 102 in 8fc4535

if len(mb.availableBufs) > 0 && mb.availableBufs[0].cap >= size {

When the kvMemBuf tries to allocate buffer, it only checks the first byte buffer's capacity in the available byte buffer pool. If the capacity is less than the allocating size, it'll create a new byte buffer, which will allocate a new pieces of memory using manual.New().
Since all the byte buffers are recycled by putting back into the available byte buffer pool, this means, each time an allocation with a size bigger than the first byte buffer's capacity, a new byte buffer will be generated with extra memory allocated, and those memory will not be freed until the chunk restore is finished.
However, when strict-mode is set to false, the large file is treated as a single chunk. The chunk restore process cannot finish in a short period of time, and the byte buffers are accumulating more and more, resulting in memory leak and potential OOM.