github.com/fiatjaf/generic-ristretto@v0.0.1/z/calloc_jemalloc.go (about)

     1  // Copyright 2020 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  // +build jemalloc
     6  
     7  package z
     8  
     9  /*
    10  #cgo LDFLAGS: /usr/local/lib/libjemalloc.a -L/usr/local/lib -Wl,-rpath,/usr/local/lib -ljemalloc -lm -lstdc++ -pthread -ldl
    11  #include <stdlib.h>
    12  #include <jemalloc/jemalloc.h>
    13  */
    14  import "C"
    15  import (
    16  	"bytes"
    17  	"fmt"
    18  	"sync"
    19  	"sync/atomic"
    20  	"unsafe"
    21  
    22  	"github.com/dustin/go-humanize"
    23  )
    24  
    25  // The go:linkname directives provides backdoor access to private functions in
    26  // the runtime. Below we're accessing the throw function.
    27  
    28  //go:linkname throw runtime.throw
    29  func throw(s string)
    30  
    31  // New allocates a slice of size n. The returned slice is from manually managed
    32  // memory and MUST be released by calling Free. Failure to do so will result in
    33  // a memory leak.
    34  //
    35  // Compile jemalloc with ./configure --with-jemalloc-prefix="je_"
    36  // https://android.googlesource.com/platform/external/jemalloc_new/+/6840b22e8e11cb68b493297a5cd757d6eaa0b406/TUNING.md
    37  // These two config options seems useful for frequent allocations and deallocations in
    38  // multi-threaded programs (like we have).
    39  // JE_MALLOC_CONF="background_thread:true,metadata_thp:auto"
    40  //
    41  // Compile Go program with `go build -tags=jemalloc` to enable this.
    42  
    43  type dalloc struct {
    44  	t  string
    45  	sz int
    46  }
    47  
    48  var dallocsMu sync.Mutex
    49  var dallocs map[unsafe.Pointer]*dalloc
    50  
    51  func init() {
    52  	// By initializing dallocs, we can start tracking allocations and deallocations via z.Calloc.
    53  	dallocs = make(map[unsafe.Pointer]*dalloc)
    54  }
    55  
    56  func Calloc(n int, tag string) []byte {
    57  	if n == 0 {
    58  		return make([]byte, 0)
    59  	}
    60  	// We need to be conscious of the Cgo pointer passing rules:
    61  	//
    62  	//   https://golang.org/cmd/cgo/#hdr-Passing_pointers
    63  	//
    64  	//   ...
    65  	//   Note: the current implementation has a bug. While Go code is permitted
    66  	//   to write nil or a C pointer (but not a Go pointer) to C memory, the
    67  	//   current implementation may sometimes cause a runtime error if the
    68  	//   contents of the C memory appear to be a Go pointer. Therefore, avoid
    69  	//   passing uninitialized C memory to Go code if the Go code is going to
    70  	//   store pointer values in it. Zero out the memory in C before passing it
    71  	//   to Go.
    72  
    73  	ptr := C.je_calloc(C.size_t(n), 1)
    74  	if ptr == nil {
    75  		// NB: throw is like panic, except it guarantees the process will be
    76  		// terminated. The call below is exactly what the Go runtime invokes when
    77  		// it cannot allocate memory.
    78  		throw("out of memory")
    79  	}
    80  
    81  	uptr := unsafe.Pointer(ptr)
    82  	dallocsMu.Lock()
    83  	dallocs[uptr] = &dalloc{
    84  		t:  tag,
    85  		sz: n,
    86  	}
    87  	dallocsMu.Unlock()
    88  	atomic.AddInt64(&numBytes, int64(n))
    89  	// Interpret the C pointer as a pointer to a Go array, then slice.
    90  	return (*[MaxArrayLen]byte)(uptr)[:n:n]
    91  }
    92  
    93  // CallocNoRef does the exact same thing as Calloc with jemalloc enabled.
    94  func CallocNoRef(n int, tag string) []byte {
    95  	return Calloc(n, tag)
    96  }
    97  
    98  // Free frees the specified slice.
    99  func Free(b []byte) {
   100  	if sz := cap(b); sz != 0 {
   101  		b = b[:cap(b)]
   102  		ptr := unsafe.Pointer(&b[0])
   103  		C.je_free(ptr)
   104  		atomic.AddInt64(&numBytes, -int64(sz))
   105  		dallocsMu.Lock()
   106  		delete(dallocs, ptr)
   107  		dallocsMu.Unlock()
   108  	}
   109  }
   110  
   111  func Leaks() string {
   112  	if dallocs == nil {
   113  		return "Leak detection disabled. Enable with 'leak' build flag."
   114  	}
   115  	dallocsMu.Lock()
   116  	defer dallocsMu.Unlock()
   117  	if len(dallocs) == 0 {
   118  		return "NO leaks found."
   119  	}
   120  	m := make(map[string]int)
   121  	for _, da := range dallocs {
   122  		m[da.t] += da.sz
   123  	}
   124  	var buf bytes.Buffer
   125  	fmt.Fprintf(&buf, "Allocations:\n")
   126  	for f, sz := range m {
   127  		fmt.Fprintf(&buf, "%s at file: %s\n", humanize.IBytes(uint64(sz)), f)
   128  	}
   129  	return buf.String()
   130  }
   131  
   132  // ReadMemStats populates stats with JE Malloc statistics.
   133  func ReadMemStats(stats *MemStats) {
   134  	if stats == nil {
   135  		return
   136  	}
   137  	// Call an epoch mallclt to refresh the stats data as mentioned in the docs.
   138  	// http://jemalloc.net/jemalloc.3.html#epoch
   139  	// Note: This epoch mallctl is as expensive as a malloc call. It takes up the
   140  	// malloc_mutex_lock.
   141  	epoch := 1
   142  	sz := unsafe.Sizeof(&epoch)
   143  	C.je_mallctl(
   144  		(C.CString)("epoch"),
   145  		unsafe.Pointer(&epoch),
   146  		(*C.size_t)(unsafe.Pointer(&sz)),
   147  		unsafe.Pointer(&epoch),
   148  		(C.size_t)(unsafe.Sizeof(epoch)))
   149  	stats.Allocated = fetchStat("stats.allocated")
   150  	stats.Active = fetchStat("stats.active")
   151  	stats.Resident = fetchStat("stats.resident")
   152  	stats.Retained = fetchStat("stats.retained")
   153  }
   154  
   155  // fetchStat is used to read a specific attribute from je malloc stats using mallctl.
   156  func fetchStat(s string) uint64 {
   157  	var out uint64
   158  	sz := unsafe.Sizeof(&out)
   159  	C.je_mallctl(
   160  		(C.CString)(s),                   // Query: eg: stats.allocated, stats.resident, etc.
   161  		unsafe.Pointer(&out),             // Variable to store the output.
   162  		(*C.size_t)(unsafe.Pointer(&sz)), // Size of the output variable.
   163  		nil,                              // Input variable used to set a value.
   164  		0) // Size of the input variable.
   165  	return out
   166  }
   167  
   168  func StatsPrint() {
   169  	opts := C.CString("mdablxe")
   170  	C.je_malloc_stats_print(nil, nil, opts)
   171  	C.free(unsafe.Pointer(opts))
   172  }