github.com/m3db/m3@v1.5.1-0.20231129193456-75a402aa583b/src/x/generics/leakcheckpool/pool.go (about)

     1  // Copyright (c) 2018 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package leakcheckpool
    22  
    23  import (
    24  	"fmt"
    25  	"runtime/debug"
    26  	"sync"
    27  	"testing"
    28  
    29  	"github.com/mauricelam/genny/generic"
    30  	"github.com/stretchr/testify/require"
    31  )
    32  
    33  // elemType is the generic type for use with the debug pool.
    34  type elemType generic.Type
    35  
    36  // elemTypeEqualsFn allows users to override equality checks
    37  // for `elemType` instances.
    38  type elemTypeEqualsFn func(a, b elemType) bool
    39  
    40  // elemTypeGetHookFn allows users to override properties on items
    41  // retrieved from the backing pools before returning in the Get()
    42  // path.
    43  type elemTypeGetHookFn func(elemType) elemType
    44  
    45  // elemTypePool is the backing pool wrapped by the debug pool.
    46  type elemTypePool interface {
    47  	generic.Type
    48  
    49  	Init()
    50  	Get() elemType
    51  	Put(elemType)
    52  }
    53  
    54  // leakcheckElemTypePoolOpts allows users to override default behaviour.
    55  type leakcheckElemTypePoolOpts struct {
    56  	DisallowUntrackedPuts bool
    57  	EqualsFn              elemTypeEqualsFn
    58  	GetHookFn             elemTypeGetHookFn
    59  }
    60  
    61  // newLeakcheckElemTypePool returns a new leakcheckElemTypePool.
    62  // nolint
    63  func newLeakcheckElemTypePool(opts leakcheckElemTypePoolOpts, backingPool elemTypePool) *leakcheckElemTypePool {
    64  	if opts.EqualsFn == nil {
    65  		// NB(prateek): fall-back to == in the worst case
    66  		opts.EqualsFn = func(a, b elemType) bool {
    67  			return a == b
    68  		}
    69  	}
    70  	return &leakcheckElemTypePool{opts: opts, elemTypePool: backingPool}
    71  }
    72  
    73  // leakcheckElemTypePool wraps the underlying elemTypePool to make it easier to
    74  // track leaks/allocs.
    75  type leakcheckElemTypePool struct {
    76  	sync.Mutex
    77  	elemTypePool
    78  	NumGets      int
    79  	NumPuts      int
    80  	PendingItems []leakcheckElemType
    81  	AllGetItems  []leakcheckElemType
    82  
    83  	opts leakcheckElemTypePoolOpts
    84  }
    85  
    86  // leakcheckElemType wraps `elemType` instances along with their last Get() paths.
    87  type leakcheckElemType struct {
    88  	Value         elemType
    89  	GetStacktrace []byte // GetStacktrace is the stacktrace for the Get() of this item
    90  }
    91  
    92  func (p *leakcheckElemTypePool) Init() {
    93  	p.Lock()
    94  	defer p.Unlock()
    95  	p.elemTypePool.Init()
    96  }
    97  
    98  func (p *leakcheckElemTypePool) Get() elemType {
    99  	p.Lock()
   100  	defer p.Unlock()
   101  
   102  	e := p.elemTypePool.Get()
   103  	if fn := p.opts.GetHookFn; fn != nil {
   104  		e = fn(e)
   105  	}
   106  
   107  	p.NumGets++
   108  	item := leakcheckElemType{
   109  		Value:         e,
   110  		GetStacktrace: debug.Stack(),
   111  	}
   112  	p.PendingItems = append(p.PendingItems, item)
   113  	p.AllGetItems = append(p.AllGetItems, item)
   114  
   115  	return e
   116  }
   117  
   118  func (p *leakcheckElemTypePool) Put(value elemType) {
   119  	p.Lock()
   120  	defer p.Unlock()
   121  
   122  	idx := -1
   123  	for i, item := range p.PendingItems {
   124  		if p.opts.EqualsFn(item.Value, value) {
   125  			idx = i
   126  			break
   127  		}
   128  	}
   129  
   130  	if idx == -1 && p.opts.DisallowUntrackedPuts {
   131  		panic(fmt.Errorf("untracked object (%v) returned to pool", value))
   132  	}
   133  
   134  	if idx != -1 {
   135  		// update slice
   136  		p.PendingItems = append(p.PendingItems[:idx], p.PendingItems[idx+1:]...)
   137  	}
   138  	p.NumPuts++
   139  
   140  	p.elemTypePool.Put(value)
   141  }
   142  
   143  // Check ensures there are no leaks.
   144  func (p *leakcheckElemTypePool) Check(t *testing.T) {
   145  	p.Lock()
   146  	defer p.Unlock()
   147  
   148  	require.Equal(t, p.NumGets, p.NumPuts)
   149  	require.Empty(t, p.PendingItems)
   150  }
   151  
   152  type leakcheckElemTypeFn func(e leakcheckElemType)
   153  
   154  // CheckExtended ensures there are no leaks, and executes the specified fn
   155  func (p *leakcheckElemTypePool) CheckExtended(t *testing.T, fn leakcheckElemTypeFn) {
   156  	p.Check(t)
   157  	p.Lock()
   158  	defer p.Unlock()
   159  	for _, e := range p.AllGetItems {
   160  		fn(e)
   161  	}
   162  }