github.com/cockroachdb/pebble@v1.1.1-0.20240513155919-3622ade60459/sstable/buffer_pool_test.go (about)

     1  // Copyright 2023 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 sstable
     6  
     7  import (
     8  	"bytes"
     9  	"fmt"
    10  	"io"
    11  	"testing"
    12  
    13  	"github.com/cockroachdb/datadriven"
    14  )
    15  
    16  func writeBufferPool(w io.Writer, bp *BufferPool) {
    17  	for i := 0; i < cap(bp.pool); i++ {
    18  		if i > 0 {
    19  			fmt.Fprint(w, " ")
    20  		}
    21  		if i >= len(bp.pool) {
    22  			fmt.Fprint(w, "[    ]")
    23  			continue
    24  		}
    25  		sz := len(bp.pool[i].v.Buf())
    26  		if bp.pool[i].b == nil {
    27  			fmt.Fprintf(w, "[%4d]", sz)
    28  		} else {
    29  			fmt.Fprintf(w, "<%4d>", sz)
    30  		}
    31  	}
    32  }
    33  
    34  func TestBufferPool(t *testing.T) {
    35  	var bp BufferPool
    36  	var buf bytes.Buffer
    37  	handles := map[string]Buf{}
    38  	drainPool := func() {
    39  		for h, b := range handles {
    40  			b.Release()
    41  			delete(handles, h)
    42  		}
    43  		bp.Release()
    44  	}
    45  	defer drainPool()
    46  	datadriven.RunTest(t, "testdata/buffer_pool", func(t *testing.T, td *datadriven.TestData) string {
    47  		buf.Reset()
    48  		switch td.Cmd {
    49  		case "init":
    50  			if cap(bp.pool) > 0 {
    51  				drainPool()
    52  			}
    53  			var initialSize int
    54  			td.ScanArgs(t, "size", &initialSize)
    55  			bp.Init(initialSize)
    56  			writeBufferPool(&buf, &bp)
    57  			return buf.String()
    58  		case "alloc":
    59  			var n int
    60  			var handle string
    61  			td.ScanArgs(t, "n", &n)
    62  			td.ScanArgs(t, "handle", &handle)
    63  			handles[handle] = bp.Alloc(n)
    64  			writeBufferPool(&buf, &bp)
    65  			return buf.String()
    66  		case "release":
    67  			var handle string
    68  			td.ScanArgs(t, "handle", &handle)
    69  			b := handles[handle]
    70  			b.Release()
    71  			delete(handles, handle)
    72  			writeBufferPool(&buf, &bp)
    73  			return buf.String()
    74  		default:
    75  			return fmt.Sprintf("unrecognized command %q", td.Cmd)
    76  		}
    77  	})
    78  }