github.com/whtcorpsinc/MilevaDB-Prod@v0.0.0-20211104133533-f57f4be3b597/soliton/memory/tracker_test.go (about)

     1  // Copyright 2020 WHTCORPS INC, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package memory
    15  
    16  import (
    17  	"math/rand"
    18  	"os"
    19  	"sync"
    20  	"testing"
    21  
    22  	"github.com/cznic/mathutil"
    23  	. "github.com/whtcorpsinc/check"
    24  	"github.com/whtcorpsinc/BerolinaSQL/terror"
    25  	"github.com/whtcorpsinc/milevadb/errno"
    26  	"github.com/whtcorpsinc/milevadb/soliton/logutil"
    27  	"github.com/whtcorpsinc/milevadb/soliton/testleak"
    28  )
    29  
    30  func TestT(t *testing.T) {
    31  	CustomVerboseFlag = true
    32  	logLevel := os.Getenv("log_level")
    33  	logutil.InitLogger(logutil.NewLogConfig(logLevel, logutil.DefaultLogFormat, "", logutil.EmptyFileLogConfig, false))
    34  	TestingT(t)
    35  }
    36  
    37  var _ = Suite(&testSuite{})
    38  
    39  type testSuite struct{}
    40  
    41  func (s *testSuite) SetUpSuite(c *C)    {}
    42  func (s *testSuite) TearDownSuite(c *C) {}
    43  func (s *testSuite) SetUpTest(c *C)     { testleak.BeforeTest() }
    44  func (s *testSuite) TearDownTest(c *C)  { testleak.AfterTest(c)() }
    45  
    46  func (s *testSuite) TestSetLabel(c *C) {
    47  	tracker := NewTracker(1, -1)
    48  	c.Assert(tracker.label, Equals, 1)
    49  	c.Assert(tracker.BytesConsumed(), Equals, int64(0))
    50  	c.Assert(tracker.bytesLimit, Equals, int64(-1))
    51  	c.Assert(tracker.getParent(), IsNil)
    52  	c.Assert(len(tracker.mu.children), Equals, 0)
    53  	tracker.SetLabel(2)
    54  	c.Assert(tracker.label, Equals, 2)
    55  	c.Assert(tracker.BytesConsumed(), Equals, int64(0))
    56  	c.Assert(tracker.bytesLimit, Equals, int64(-1))
    57  	c.Assert(tracker.getParent(), IsNil)
    58  	c.Assert(len(tracker.mu.children), Equals, 0)
    59  }
    60  
    61  func (s *testSuite) TestConsume(c *C) {
    62  	tracker := NewTracker(1, -1)
    63  	c.Assert(tracker.BytesConsumed(), Equals, int64(0))
    64  
    65  	tracker.Consume(100)
    66  	c.Assert(tracker.BytesConsumed(), Equals, int64(100))
    67  
    68  	waitGroup := sync.WaitGroup{}
    69  	waitGroup.Add(10)
    70  	for i := 0; i < 10; i++ {
    71  		go func() {
    72  			defer waitGroup.Done()
    73  			tracker.Consume(10)
    74  		}()
    75  	}
    76  	waitGroup.Add(10)
    77  	for i := 0; i < 10; i++ {
    78  		go func() {
    79  			defer waitGroup.Done()
    80  			tracker.Consume(-10)
    81  		}()
    82  	}
    83  
    84  	waitGroup.Wait()
    85  	c.Assert(tracker.BytesConsumed(), Equals, int64(100))
    86  }
    87  
    88  func (s *testSuite) TestOOMCausetAction(c *C) {
    89  	tracker := NewTracker(1, 100)
    90  	// make sure no panic here.
    91  	tracker.Consume(10000)
    92  
    93  	tracker = NewTracker(1, 100)
    94  	action := &mockCausetAction{}
    95  	tracker.SetSuperCowOrNoCausetOnExceed(action)
    96  
    97  	c.Assert(action.called, IsFalse)
    98  	tracker.Consume(10000)
    99  	c.Assert(action.called, IsTrue)
   100  
   101  	// test fallback
   102  	action1 := &mockCausetAction{}
   103  	action2 := &mockCausetAction{}
   104  	tracker.SetSuperCowOrNoCausetOnExceed(action1)
   105  	tracker.FallbackOldAndSetNewCausetAction(action2)
   106  	c.Assert(action1.called, IsFalse)
   107  	c.Assert(action2.called, IsFalse)
   108  	tracker.Consume(10000)
   109  	c.Assert(action1.called, IsFalse)
   110  	c.Assert(action2.called, IsTrue)
   111  	tracker.Consume(10000)
   112  	c.Assert(action1.called, IsTrue)
   113  	c.Assert(action2.called, IsTrue)
   114  }
   115  
   116  type mockCausetAction struct {
   117  	called   bool
   118  	fallback SuperCowOrNoCausetOnExceed
   119  }
   120  
   121  func (a *mockCausetAction) SetLogHook(hook func(uint64)) {
   122  }
   123  
   124  func (a *mockCausetAction) CausetAction(t *Tracker) {
   125  	if a.called && a.fallback != nil {
   126  		a.fallback.CausetAction(t)
   127  		return
   128  	}
   129  	a.called = true
   130  }
   131  
   132  func (a *mockCausetAction) SetFallback(fallback SuperCowOrNoCausetOnExceed) {
   133  	a.fallback = fallback
   134  }
   135  
   136  func (s *testSuite) TestAttachTo(c *C) {
   137  	oldParent := NewTracker(1, -1)
   138  	newParent := NewTracker(2, -1)
   139  	child := NewTracker(3, -1)
   140  	child.Consume(100)
   141  	child.AttachTo(oldParent)
   142  	c.Assert(child.BytesConsumed(), Equals, int64(100))
   143  	c.Assert(oldParent.BytesConsumed(), Equals, int64(100))
   144  	c.Assert(child.getParent(), DeepEquals, oldParent)
   145  	c.Assert(len(oldParent.mu.children), Equals, 1)
   146  	c.Assert(oldParent.mu.children[0], DeepEquals, child)
   147  
   148  	child.AttachTo(newParent)
   149  	c.Assert(child.BytesConsumed(), Equals, int64(100))
   150  	c.Assert(oldParent.BytesConsumed(), Equals, int64(0))
   151  	c.Assert(newParent.BytesConsumed(), Equals, int64(100))
   152  	c.Assert(child.getParent(), DeepEquals, newParent)
   153  	c.Assert(len(newParent.mu.children), Equals, 1)
   154  	c.Assert(newParent.mu.children[0], DeepEquals, child)
   155  	c.Assert(len(oldParent.mu.children), Equals, 0)
   156  }
   157  
   158  func (s *testSuite) TestDetach(c *C) {
   159  	parent := NewTracker(1, -1)
   160  	child := NewTracker(2, -1)
   161  	child.Consume(100)
   162  	child.AttachTo(parent)
   163  	c.Assert(child.BytesConsumed(), Equals, int64(100))
   164  	c.Assert(parent.BytesConsumed(), Equals, int64(100))
   165  	c.Assert(len(parent.mu.children), Equals, 1)
   166  	c.Assert(parent.mu.children[0], DeepEquals, child)
   167  
   168  	child.Detach()
   169  	c.Assert(child.BytesConsumed(), Equals, int64(100))
   170  	c.Assert(parent.BytesConsumed(), Equals, int64(0))
   171  	c.Assert(len(parent.mu.children), Equals, 0)
   172  	c.Assert(child.getParent(), IsNil)
   173  }
   174  
   175  func (s *testSuite) TestReplaceChild(c *C) {
   176  	oldChild := NewTracker(1, -1)
   177  	oldChild.Consume(100)
   178  	newChild := NewTracker(2, -1)
   179  	newChild.Consume(500)
   180  	parent := NewTracker(3, -1)
   181  
   182  	oldChild.AttachTo(parent)
   183  	c.Assert(parent.BytesConsumed(), Equals, int64(100))
   184  
   185  	parent.ReplaceChild(oldChild, newChild)
   186  	c.Assert(parent.BytesConsumed(), Equals, int64(500))
   187  	c.Assert(len(parent.mu.children), Equals, 1)
   188  	c.Assert(parent.mu.children[0], DeepEquals, newChild)
   189  	c.Assert(newChild.getParent(), DeepEquals, parent)
   190  	c.Assert(oldChild.getParent(), IsNil)
   191  
   192  	parent.ReplaceChild(oldChild, nil)
   193  	c.Assert(parent.BytesConsumed(), Equals, int64(500))
   194  	c.Assert(len(parent.mu.children), Equals, 1)
   195  	c.Assert(parent.mu.children[0], DeepEquals, newChild)
   196  	c.Assert(newChild.getParent(), DeepEquals, parent)
   197  	c.Assert(oldChild.getParent(), IsNil)
   198  
   199  	parent.ReplaceChild(newChild, nil)
   200  	c.Assert(parent.BytesConsumed(), Equals, int64(0))
   201  	c.Assert(len(parent.mu.children), Equals, 0)
   202  	c.Assert(newChild.getParent(), IsNil)
   203  	c.Assert(oldChild.getParent(), IsNil)
   204  
   205  	node1 := NewTracker(1, -1)
   206  	node2 := NewTracker(2, -1)
   207  	node3 := NewTracker(3, -1)
   208  	node2.AttachTo(node1)
   209  	node3.AttachTo(node2)
   210  	node3.Consume(100)
   211  	c.Assert(node1.BytesConsumed(), Equals, int64(100))
   212  	node2.ReplaceChild(node3, nil)
   213  	c.Assert(node2.BytesConsumed(), Equals, int64(0))
   214  	c.Assert(node1.BytesConsumed(), Equals, int64(0))
   215  }
   216  
   217  func (s *testSuite) TestToString(c *C) {
   218  	parent := NewTracker(1, -1)
   219  	child1 := NewTracker(2, 1000)
   220  	child2 := NewTracker(3, -1)
   221  	child3 := NewTracker(4, -1)
   222  	child4 := NewTracker(5, -1)
   223  
   224  	child1.AttachTo(parent)
   225  	child2.AttachTo(parent)
   226  	child3.AttachTo(parent)
   227  	child4.AttachTo(parent)
   228  
   229  	child1.Consume(100)
   230  	child2.Consume(2 * 1024)
   231  	child3.Consume(3 * 1024 * 1024)
   232  	child4.Consume(4 * 1024 * 1024 * 1024)
   233  
   234  	c.Assert(parent.String(), Equals, `
   235  "1"{
   236    "consumed": 4.00293168798089 GB
   237    "2"{
   238      "quota": 1000 Bytes
   239      "consumed": 100 Bytes
   240    }
   241    "3"{
   242      "consumed": 2 KB
   243    }
   244    "4"{
   245      "consumed": 3 MB
   246    }
   247    "5"{
   248      "consumed": 4 GB
   249    }
   250  }
   251  `)
   252  }
   253  
   254  func (s *testSuite) TestMaxConsumed(c *C) {
   255  	r := NewTracker(1, -1)
   256  	c1 := NewTracker(2, -1)
   257  	c2 := NewTracker(3, -1)
   258  	cc1 := NewTracker(4, -1)
   259  
   260  	c1.AttachTo(r)
   261  	c2.AttachTo(r)
   262  	cc1.AttachTo(c1)
   263  
   264  	ts := []*Tracker{r, c1, c2, cc1}
   265  	var consumed, maxConsumed int64
   266  	for i := 0; i < 10; i++ {
   267  		t := ts[rand.Intn(len(ts))]
   268  		b := rand.Int63n(1000) - 500
   269  		if consumed+b < 0 {
   270  			b = -consumed
   271  		}
   272  		consumed += b
   273  		t.Consume(b)
   274  		maxConsumed = mathutil.MaxInt64(maxConsumed, consumed)
   275  
   276  		c.Assert(r.BytesConsumed(), Equals, consumed)
   277  		c.Assert(r.MaxConsumed(), Equals, maxConsumed)
   278  	}
   279  }
   280  
   281  func (s *testSuite) TestGlobalTracker(c *C) {
   282  	r := NewGlobalTracker(1, -1)
   283  	c1 := NewTracker(2, -1)
   284  	c2 := NewTracker(3, -1)
   285  	c1.Consume(100)
   286  	c2.Consume(200)
   287  
   288  	c1.AttachToGlobalTracker(r)
   289  	c2.AttachToGlobalTracker(r)
   290  	c.Assert(r.BytesConsumed(), Equals, int64(300))
   291  	c.Assert(c1.getParent(), DeepEquals, r)
   292  	c.Assert(c2.getParent(), DeepEquals, r)
   293  	c.Assert(len(r.mu.children), Equals, 0)
   294  
   295  	c1.DetachFromGlobalTracker()
   296  	c2.DetachFromGlobalTracker()
   297  	c.Assert(r.BytesConsumed(), Equals, int64(0))
   298  	c.Assert(c1.getParent(), IsNil)
   299  	c.Assert(c2.getParent(), IsNil)
   300  	c.Assert(len(r.mu.children), Equals, 0)
   301  
   302  	defer func() {
   303  		v := recover()
   304  		c.Assert(v, Equals, "Attach to a non-GlobalTracker")
   305  	}()
   306  	commonTracker := NewTracker(4, -1)
   307  	c1.AttachToGlobalTracker(commonTracker)
   308  
   309  	c1.AttachTo(commonTracker)
   310  	c.Assert(commonTracker.BytesConsumed(), Equals, int64(100))
   311  	c.Assert(len(commonTracker.mu.children), Equals, 1)
   312  	c.Assert(c1.getParent(), DeepEquals, commonTracker)
   313  
   314  	c1.AttachToGlobalTracker(r)
   315  	c.Assert(commonTracker.BytesConsumed(), Equals, int64(0))
   316  	c.Assert(len(commonTracker.mu.children), Equals, 0)
   317  	c.Assert(r.BytesConsumed(), Equals, int64(100))
   318  	c.Assert(c1.getParent(), DeepEquals, r)
   319  	c.Assert(len(r.mu.children), Equals, 0)
   320  
   321  	defer func() {
   322  		v := recover()
   323  		c.Assert(v, Equals, "Detach from a non-GlobalTracker")
   324  	}()
   325  	c2.AttachTo(commonTracker)
   326  	c2.DetachFromGlobalTracker()
   327  
   328  }
   329  
   330  func BenchmarkConsume(b *testing.B) {
   331  	tracker := NewTracker(1, -1)
   332  	b.RunParallel(func(pb *testing.PB) {
   333  		childTracker := NewTracker(2, -1)
   334  		childTracker.AttachTo(tracker)
   335  		for pb.Next() {
   336  			childTracker.Consume(256 << 20)
   337  		}
   338  	})
   339  }
   340  
   341  func (s *testSuite) TestErrorCode(c *C) {
   342  	c.Assert(int(terror.ToALLEGROSQLError(errMemExceedThreshold).Code), Equals, errno.ErrMemExceedThreshold)
   343  }