github.com/fraugster/parquet-go@v0.12.0/alloc.go (about)

     1  package goparquet
     2  
     3  import (
     4  	"fmt"
     5  	"reflect"
     6  	"runtime"
     7  	"sync"
     8  )
     9  
    10  type allocTracker struct {
    11  	mtx       sync.RWMutex
    12  	allocs    map[uintptr]uint64
    13  	totalSize uint64
    14  	maxSize   uint64
    15  }
    16  
    17  func newAllocTracker(maxSize uint64) *allocTracker {
    18  	return &allocTracker{
    19  		allocs:  make(map[uintptr]uint64),
    20  		maxSize: maxSize,
    21  	}
    22  }
    23  
    24  func (t *allocTracker) register(obj interface{}, size uint64) {
    25  	if t == nil {
    26  		return
    27  	}
    28  
    29  	t.mtx.Lock()
    30  	defer t.mtx.Unlock()
    31  
    32  	if _, ok := obj.([]byte); ok {
    33  		obj = &obj
    34  	}
    35  
    36  	key := reflect.ValueOf(obj).Pointer()
    37  
    38  	if _, ok := t.allocs[key]; ok { // object has already been tracked, no need to add it.
    39  		t.mtx.Unlock()
    40  		return
    41  	}
    42  
    43  	t.allocs[key] = size
    44  	t.totalSize += size
    45  
    46  	runtime.SetFinalizer(obj, t.finalize)
    47  
    48  	if t.maxSize > 0 && t.totalSize > t.maxSize {
    49  		t.doPanic(t.totalSize)
    50  	}
    51  }
    52  
    53  func (t *allocTracker) test(size uint64) {
    54  	if t == nil {
    55  		return
    56  	}
    57  	t.mtx.RLock()
    58  	defer t.mtx.RUnlock()
    59  	if t.maxSize > 0 && t.totalSize+size > t.maxSize {
    60  		t.doPanic(t.totalSize + size)
    61  	}
    62  }
    63  
    64  func (t *allocTracker) doPanic(totalSize uint64) {
    65  	if t == nil {
    66  		return
    67  	}
    68  	panic(fmt.Errorf("memory usage of %d bytes is greater than configured maximum of %d bytes", totalSize, t.maxSize))
    69  }
    70  
    71  func (t *allocTracker) finalize(obj interface{}) {
    72  	if t == nil {
    73  		return
    74  	}
    75  
    76  	t.mtx.Lock()
    77  	defer t.mtx.Unlock()
    78  
    79  	key := reflect.ValueOf(obj).Pointer()
    80  
    81  	size, ok := t.allocs[key]
    82  	if !ok { // if object hasn't been tracked, do nothing.
    83  		return
    84  	}
    85  
    86  	// remove size from total size, and unregister from tracker.
    87  	t.totalSize -= size
    88  	delete(t.allocs, key)
    89  }