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  }