diff --git a/util/memory/tracker.go b/util/memory/tracker.go index eaf8e3d0cbbea..54ca1d75fa255 100644 --- a/util/memory/tracker.go +++ b/util/memory/tracker.go @@ -217,7 +217,12 @@ func reArrangeFallback(a ActionOnExceed, b ActionOnExceed) ActionOnExceed { // SetLabel sets the label of a Tracker. func (t *Tracker) SetLabel(label int) { + parent := t.getParent() + t.Detach() t.label = label + if parent != nil { + t.AttachTo(parent) + } } // Label gets the label of a Tracker. @@ -229,6 +234,10 @@ func (t *Tracker) Label() int { // already has a parent, this function will remove it from the old parent. // Its consumed memory usage is used to update all its ancestors. func (t *Tracker) AttachTo(parent *Tracker) { + if parent.isGlobal { + t.AttachToGlobalTracker(parent) + return + } oldParent := t.getParent() if oldParent != nil { oldParent.remove(t) diff --git a/util/memory/tracker_test.go b/util/memory/tracker_test.go index d82a8364014b4..263de711b5da2 100644 --- a/util/memory/tracker_test.go +++ b/util/memory/tracker_test.go @@ -43,6 +43,18 @@ func TestSetLabel(t *testing.T) { require.Equal(t, 0, len(tracker.mu.children)) } +func TestSetLabel2(t *testing.T) { + tracker := NewTracker(1, -1) + tracker2 := NewTracker(2, -1) + tracker2.AttachTo(tracker) + tracker2.Consume(10) + require.Equal(t, tracker.BytesConsumed(), int64(10)) + tracker2.SetLabel(10) + require.Equal(t, tracker.BytesConsumed(), int64(10)) + tracker2.Detach() + require.Equal(t, tracker.BytesConsumed(), int64(0)) +} + func TestConsume(t *testing.T) { tracker := NewTracker(1, -1) require.Equal(t, int64(0), tracker.BytesConsumed())