Repository: pingcap/tidb
File: util/memory/tracker.go (Go)

Code before:
// distinguish between "0 bytes" and "N/A". ref: binaryOpFromFlatOp()
func (t *Tracker) MaxConsumed() int64 {
	return t.maxConsumed.Load()
}

// SearchTrackerWithoutLock searches the specific tracker under this tracker without lock.
func (t *Tracker) SearchTrackerWithoutLock(label int) *Tracker {
	if t.label == label {
		return t
	}
	children := t.mu.children[label]
	if len(children) > 0 {
		return children[0]
	}
	return nil
}

// SearchTrackerConsumedMoreThanNBytes searches the specific tracker that consumes more than NBytes.
func (t *Tracker) SearchTrackerConsumedMoreThanNBytes(limit int64) (res []*Tracker) {
	t.mu.Lock()
	defer t.mu.Unlock()
	for _, childSlice := range t.mu.children {
		for _, tracker := range childSlice {
			if tracker.BytesConsumed() > limit {
				res = append(res, tracker)
			}
		}
	}
	return
}

// String returns the string representation of this Tracker tree.
func (t *Tracker) String() string {
	buffer := bytes.NewBufferString("\n")
	t.toString("", buffer)
	return buffer.String()
}

func (t *Tracker) toString(indent string, buffer *bytes.Buffer) {
	fmt.Fprintf(buffer, "%s\"%d\"{\n", indent, t.label)
	bytesLimit := t.GetBytesLimit()
	if bytesLimit > 0 {
		fmt.Fprintf(buffer, "%s  \"quota\": %s\n", indent, t.FormatBytes(bytesLimit))
	}
	fmt.Fprintf(buffer, "%s  \"consumed\": %s\n", indent, t.FormatBytes(t.BytesConsumed()))

	t.mu.Lock()
	labels := make([]int, 0, len(t.mu.children))
	for label := range t.mu.children {
		labels = append(labels, label)
	}

Review comment (performance):
The Release method creates a finalizer for every release operation when GC-aware tracking is enabled. This can create a memory leak if the same bytes value is used multiple times in the finalizer closure. The finalizer captures the `bytes` and `tracker` variables, but multiple finalizers might be registered with the same tracker before GC runs. If bytes accumulate faster than GC runs, this could lead to memory pressure from the finalizer queue itself.

Consider tracking a single finalizer per tracker instead of creating one per Release call, or implement a batching mechanism to reduce the number of finalizers created.

Code after:
	if *bufferedMemSize >= int64(TrackMemWhenExceeds) {
		t.Consume(*bufferedMemSize)
		*bufferedMemSize = int64(0)
	}
}

// Release is used to release memory tracked, track the released memory until GC triggered if needed
// If you want your track to be GC-aware, please use Release(bytes) instead of Consume(-bytes), and pass the memory size of the real object.
// Only Analyze is integrated with Release so far.
func (t *Tracker) Release(bytes int64) {
	if bytes == 0 {
		return
	}
	defer t.Consume(-bytes)
	for tracker := t; tracker != nil; tracker = tracker.getParent() {
		if tracker.shouldRecordRelease() {
			// use fake ref instead of obj ref, otherwise obj will be reachable again and gc in next cycle
			newRef := &finalizerRef{}
			runtime.SetFinalizer(newRef, func(ref *finalizerRef) {
				tracker.release(bytes)
			})
			tracker.recordRelease(bytes)
			return
		}
	}
}

// BufferedRelease is used to buffer memory release and do late release
// not thread-safe, should be called in one goroutine
func (t *Tracker) BufferedRelease(bufferedMemSize *int64, bytes int64) {
	*bufferedMemSize += bytes
	if *bufferedMemSize >= int64(TrackMemWhenExceeds) {
		t.Release(*bufferedMemSize)
		*bufferedMemSize = int64(0)
	}
}

func (t *Tracker) shouldRecordRelease() bool {
	return EnableGCAwareMemoryTrack.Load() && t.label == LabelForGlobalAnalyzeMemory
}

func (t *Tracker) recordRelease(bytes int64) {
	for tracker := t; tracker != nil; tracker = tracker.getParent() {
		bytesReleased := atomic.AddInt64(&tracker.bytesReleased, bytes)
		if label, ok := MetricsTypes[tracker.label]; ok {
			metrics.MemoryUsage.WithLabelValues(label[0], label[2]).Set(float64(bytesReleased))
		}
	}
}

func (t *Tracker) release(bytes int64) {