unlock unlimited code runs for $79, limited time →
NOSSLETTER.techLIVE
back to OSSpedia
pingcap/tidb
mid7/10
memory-managementerror-handlingtesting-strategy

server: limit the size for COM_STMT_SEND_LONG_DATA in each connection (release-7.5-20260724-v7.5.7) (#70350)

Project: pingcap/tidb | Language: Go | Concepts: memory-management, error-handling, testing-strategy | Difficulty: mid
Link: https://github.com/pingcap/tidb/pull/70442

The problem

Prepared statements in TiDB allow clients to send large parameter data in chunks using the COM_STMT_SEND_LONG_DATA protocol command. Previously, there was no effective per-connection memory limit on how much data could be accumulated this way before statement execution. This could allow a client to consume excessive memory, potentially exceeding the configured per-query memory quota (tidb_mem_quota_query) and destabilizing the server.

Why it's non-trivial

A naive fix might just check the size of each incoming chunk or the total parameter size when executing the statement. However, this would not prevent the server from accumulating unbounded memory usage before the statement is executed. Additionally, memory must be released correctly when the statement is reset or closed, and error handling needs to distinguish between different reasons for refusing more data (e.g., packet size vs. memory quota).

The approach

This PR introduces explicit memory tracking and enforcement for long-data parameters at the statement level, integrates with the session's memory tracker, and ensures proper cleanup. Here are the key changes:

1. Track and enforce memory usage for long-data parameters

A new field is added to the statement struct to track bytes held by long-data parameters, and a flag to record if the memory quota was exceeded:

+	// boundParamsMemQuotaExceeded tracks whether further COM_STMT_SEND_LONG_DATA
+	// was refused because the session MemTracker would exceed tidb_mem_quota_query.
+	boundParamsMemQuotaExceeded bool
+	// boundLongDataBytes is the number of COM_STMT_SEND_LONG_DATA bytes currently
+	// held by this statement and charged to the session MemTracker.
+	boundLongDataBytes int64

What this does:
These fields allow the statement to know both how much memory it has charged to the session's memory tracker, and whether it has already refused further data due to exceeding the per-query quota.

2. Charge and release memory on parameter changes

When appending parameter data, the code checks the quota before consuming memory, and releases memory when data is reset or closed:

+	// Once either limit has been exceeded for this statement, keep SEND_LONG_DATA
+	// silent and refuse further growth until RESET/EXECUTE clears the buffers.
+	if ts.boundParamsTooLarge || ts.boundParamsMemQuotaExceeded {
+		return nil
+	}
+
+	chunkSize := int64(len(data))
+	vars := ts.ctx.GetSessionVars()
+	if uint64(len(ts.boundParams[paramID]))+uint64(chunkSize) > vars.MaxAllowedPacket {
+		// MySQL reports the packet-too-large error on the following EXECUTE, not on SEND_LONG_DATA.
+		// Stop appending more bytes once the limit is exceeded so the statement cannot grow unboundedly.
+		ts.boundParamsTooLarge = true
+		return nil
+	}
+
+	// Charge COM_STMT_SEND_LONG_DATA to the session MemTracker so it counts toward
+	// tidb_mem_quota_query. Refuse before Consume to avoid triggering OOM actions
+	// on a protocol command that has no response packet.
+	memTracker := vars.MemTracker
+	quota := vars.MemQuotaQuery
+	if quota > 0 && memTracker.BytesConsumed()+chunkSize >= quota {
+		ts.boundParamsMemQuotaExceeded = true
+		return nil
+	}
+
+	memTracker.Consume(chunkSize)
+	ts.boundParams[paramID] = append(ts.boundParams[paramID], data...)
+	ts.boundLongDataBytes += chunkSize

What this does:

  • If the quota would be exceeded by this chunk, refuse the data and set a flag.
  • If the data is accepted, charge the memory to the session's tracker.
  • If the data is empty, release any previously charged memory for that parameter.

3. Release memory on statement reset and close

+func (ts *TiDBStatement) releaseBoundLongData() {
+	if ts.boundLongDataBytes > 0 {
+		ts.ctx.GetSessionVars().MemTracker.Consume(-ts.boundLongDataBytes)
+		ts.boundLongDataBytes = 0
+	}
+	for i := range ts.boundParams {
+		ts.boundParams[i] = nil
+	}
+	ts.boundParamsTooLarge = false
+	ts.boundParamsMemQuotaExceeded = false
+}

This function is called from both Reset and Close:

// Reset implements PreparedStatement Reset method.
func (ts *TiDBStatement) Reset() error {
-	for i := range ts.boundParams {
-		ts.boundParams[i] = nil
-	}
-	ts.boundParamsTooLarge = false
+	ts.releaseBoundLongData()
...
}

// Close implements PreparedStatement Close method.
func (ts *TiDBStatement) Close() error {
+	ts.releaseBoundLongData()
...
}

What this does:
Ensures that all memory charged for long-data parameters is released when the statement is reset or closed, and that the state flags are cleared.

4. Distinguish between memory quota and packet-size errors

When checking if the statement is ready for execution, the code now reports a specific error if the memory quota was exceeded:

func (ts *TiDBStatement) CheckLongDataSize() error {
+	if ts.boundParamsMemQuotaExceeded {
+		return errors.Errorf("%s%s[conn=%d]", memory.PanicMemoryExceedWarnMsg,
+			memory.WarnMsgSuffixForSingleQuery, ts.ctx.GetSessionVars().ConnectionID)
+	}
	if ts.boundParamsTooLarge {
		return servererr.ErrNetPacketTooLarge
	}
	return nil
}

What this does:
Allows the client to distinguish between exceeding max_allowed_packet and exceeding the per-query memory quota.

5. Add comprehensive tests for quota enforcement and cleanup

A new unit test exercises the memory quota logic, including refusing data, releasing memory, and accepting new data after cleanup:

func TestStmtSendLongDataMemQuotaQuery(t *testing.T) {
	...
	require.NoError(t, dispatchSendLongData(c, stmt.ID(), 0, bytes.Repeat([]byte{'a'}, 600)))
	require.Equal(t, base+600, vars.MemTracker.BytesConsumed())
	...
	require.NoError(t, dispatchSendLongData(c, stmt.ID(), 0, bytes.Repeat([]byte{'b'}, 500)))
	...
	err = c.Dispatch(context.Background(), ...)
	require.Error(t, err)
	require.Contains(t, err.Error(), memory.PanicMemoryExceedWarnMsg+memory.WarnMsgSuffixForSingleQuery)
	require.Equal(t, base, vars.MemTracker.BytesConsumed())
	...
	require.NoError(t, dispatchSendLongData(c, stmt.ID(), 0, bytes.Repeat([]byte{'c'}, 1023)))
	require.Len(t, stmt.BoundParams()[0], 1023)
	require.Equal(t, base+1023, vars.MemTracker.BytesConsumed())
	require.NoError(t, stmt.Close())
	require.Equal(t, base, vars.MemTracker.BytesConsumed())
}

What this does:

  • Sends enough data to approach the quota, confirms memory is charged.
  • Attempts to exceed the quota, expects refusal and no further memory growth.
  • Checks that memory is released after execution or close, and that new data can be accepted after cleanup.

From the review thread

No notable review discussion on this PR.

Takeaways

  • Memory for protocol-level buffers should be tracked and limited at the earliest opportunity, not just at execution time.
  • Resource cleanup must be tied to all relevant lifecycle events (reset, close, error) to prevent memory leaks or quota lockout.
  • Error handling should distinguish between different resource limits (e.g., packet size vs. memory quota) to provide actionable feedback to clients.