github.com/petermattis/pebble@v0.0.0-20190905164901-ab51a2166067/pacer_test.go (about) 1 // Copyright 2019 The LevelDB-Go and Pebble Authors. All rights reserved. Use 2 // of this source code is governed by a BSD-style license that can be found in 3 // the LICENSE file. 4 5 package pebble 6 7 import ( 8 "bytes" 9 "context" 10 "fmt" 11 "strconv" 12 "strings" 13 "testing" 14 "time" 15 16 "github.com/petermattis/pebble/internal/datadriven" 17 ) 18 19 type mockCountLimiter struct { 20 waitCount int 21 allowCount int 22 burst int 23 } 24 25 func (m *mockCountLimiter) WaitN(ctx context.Context, n int) error { 26 m.waitCount += n 27 return nil 28 } 29 30 func (m *mockCountLimiter) AllowN(now time.Time, n int) bool { 31 m.allowCount += n 32 return true 33 } 34 35 func (m *mockCountLimiter) Burst() int { 36 return m.burst 37 } 38 39 type mockPrintLimiter struct { 40 buf bytes.Buffer 41 burst int 42 } 43 44 func (m *mockPrintLimiter) WaitN(ctx context.Context, n int) error { 45 fmt.Fprintf(&m.buf, "wait: %d\n", n) 46 return nil 47 } 48 49 func (m *mockPrintLimiter) AllowN(now time.Time, n int) bool { 50 fmt.Fprintf(&m.buf, "allow: %d\n", n) 51 return true 52 } 53 54 func (m *mockPrintLimiter) Burst() int { 55 return m.burst 56 } 57 58 func TestCompactionPacerMaybeThrottle(t *testing.T) { 59 datadriven.RunTest(t, "testdata/compaction_pacer_maybe_throttle", 60 func(d *datadriven.TestData) string { 61 switch d.Cmd { 62 case "init": 63 if len(d.CmdArgs) != 1 { 64 return fmt.Sprintf("%s expects 1 argument", d.Cmd) 65 } 66 67 burst := uint64(1) 68 dirtyBytes := uint64(1) 69 var bytesIterated uint64 70 var currentTotal uint64 71 var slowdownThreshold uint64 72 if len(d.Input) > 0 { 73 for _, data := range strings.Split(d.Input, "\n") { 74 parts := strings.Split(data, ":") 75 if len(parts) != 2 { 76 return fmt.Sprintf("malformed test:\n%s", d.Input) 77 } 78 varKey := parts[0] 79 varValue, err := strconv.ParseUint(strings.TrimSpace(parts[1]), 10, 64) 80 if err != nil { 81 return err.Error() 82 } 83 84 switch varKey { 85 case "burst": 86 burst = varValue 87 case "bytesIterated": 88 bytesIterated = varValue 89 case "currentTotal": 90 currentTotal = varValue 91 case "dirtyBytes": 92 dirtyBytes = varValue 93 case "slowdownThreshold": 94 slowdownThreshold = varValue 95 default: 96 return fmt.Sprintf("unknown command: %s", varKey) 97 } 98 } 99 } 100 101 mockLimiter := mockPrintLimiter{burst: int(burst)} 102 switch d.CmdArgs[0].Key { 103 case "compaction": 104 getInfo := func() compactionPacerInfo { 105 return compactionPacerInfo{ 106 slowdownThreshold: slowdownThreshold, 107 totalCompactionDebt: currentTotal, 108 totalDirtyBytes: dirtyBytes, 109 } 110 } 111 compactionPacer := newCompactionPacer(compactionPacerEnv{ 112 limiter: &mockLimiter, 113 memTableSize: 100, 114 getInfo: getInfo, 115 }) 116 117 err := compactionPacer.maybeThrottle(bytesIterated) 118 if err != nil { 119 return err.Error() 120 } 121 122 return mockLimiter.buf.String() 123 case "flush": 124 getInfo := func() flushPacerInfo { 125 return flushPacerInfo{ 126 totalBytes: currentTotal, 127 } 128 } 129 flushPacer := newFlushPacer(flushPacerEnv{ 130 limiter: &mockLimiter, 131 memTableSize: 100, 132 getInfo: getInfo, 133 }) 134 flushPacer.slowdownThreshold = slowdownThreshold 135 136 err := flushPacer.maybeThrottle(bytesIterated) 137 if err != nil { 138 return err.Error() 139 } 140 141 return mockLimiter.buf.String() 142 default: 143 return fmt.Sprintf("unknown command: %s", d.Cmd) 144 } 145 146 default: 147 return fmt.Sprintf("unknown command: %s", d.Cmd) 148 } 149 }) 150 }