Mitigating Memory Fragmentation in High-Throughput Go Services with Custom jemalloc Allocators — code snippets

snippet-1.go

package jemalloc

/*
#cgo LDFLAGS: -ljemalloc
#include <stdlib.h>
#include <jemalloc/jemalloc.h>
*/
import "C"
import (
	"unsafe"
)

// Alloc allocates size bytes off-heap using jemalloc.
func Alloc(size int) unsafe.Pointer {
	ptr := C.je_malloc(C.size_t(size))
	return ptr
}

// Free releases the off-heap memory allocated by Alloc.
func Free(ptr unsafe.Pointer) {
	C.je_free(ptr)
}

// Realloc resizes the allocated memory block to newSize.
func Realloc(ptr unsafe.Pointer, newSize int) unsafe.Pointer {
	return C.je_realloc(ptr, C.size_t(newSize))
}

snippet-2.go

package jemalloc

import (
	"unsafe"
)

// AllocSlice allocates a byte slice of the specified capacity directly from jemalloc.
// The returned slice has len == cap and cap == cap. The memory is not zero-initialized.
func AllocSlice(cap int) []byte {
	if cap <= 0 {
		return nil
	}
	ptr := Alloc(cap)
	if ptr == nil {
		panic("jemalloc: out of memory")
	}
	// Return a slice pointing to the jemalloc allocated memory block.
	// unsafe.Slice is available in Go 1.17+ and is the standard way to create slices from raw pointers.
	return unsafe.Slice((*byte)(ptr), cap)
}

// FreeSlice extracts the underlying pointer of the slice and frees it via jemalloc.
// The slice must not be read or written to after calling this function.
func FreeSlice(slice []byte) {
	if cap(slice) == 0 {
		return
	}
	// unsafe.SliceData retrieves the pointer to the underlying array (Go 1.20+).
	ptr := unsafe.Pointer(unsafe.SliceData(slice))
	Free(ptr)
}

snippet-3.go

package jemalloc

import (
	"runtime"
	"sync/atomic"
	"unsafe"
)

// OffHeapBuffer wraps a jemalloc-allocated byte slice.
// It is designed to prevent double-free issues and detect leaks.
type OffHeapBuffer struct {
	data  []byte
	freed int32
}

// NewBuffer allocates an OffHeapBuffer of the specified size.
func NewBuffer(size int) *OffHeapBuffer {
	buf := &OffHeapBuffer{
		data: AllocSlice(size),
	}
	
	// Register a finalizer as an emergency fallback.
	// If the developer forgets to call Release(), the Go GC will eventually clean it up.
	// WARNING: Relying on finalizers is not a substitute for explicit Release calls.
	runtime.SetFinalizer(buf, func(b *OffHeapBuffer) {
		b.Release()
	})
	
	return buf
}

// Bytes returns the underlying byte slice pointing to off-heap memory.
func (b *OffHeapBuffer) Bytes() []byte {
	if atomic.LoadInt32(&b.freed) == 1 {
		panic("use of freed jemalloc buffer")
	}
	return b.data
}

// Release returns the memory back to jemalloc and removes the finalizer.
func (b *OffHeapBuffer) Release() {
	if atomic.CompareAndSwapInt32(&b.freed, 0, 1) {
		FreeSlice(b.data)
		b.data = nil
		// Clear the finalizer so the garbage collector can free the struct itself.
		runtime.SetFinalizer(b, nil)
	}
}

snippet-4.go

package jemalloc

/*
#cgo LDFLAGS: -ljemalloc
#include <stdlib.h>
#include <jemalloc/jemalloc.h>
*/
import "C"
import (
	"fmt"
	"unsafe"
)

// SetDecayTimes configures jemalloc's decay times globally for all future arenas.
// dirtyDecayMS: time in milliseconds before reclaiming dirty pages.
// muzzyDecayMS: time in milliseconds before reclaiming muzzy pages.
func SetDecayTimes(dirtyDecayMS, muzzyDecayMS int) error {
	cDirty := C.ssize_t(dirtyDecayMS)
	cMuzzy := C.ssize_t(muzzyDecayMS)

	// Update arenas.dirty_decay_ms
	dirtyKey := C.CString("arenas.dirty_decay_ms")
	defer C.free(unsafe.Pointer(dirtyKey))

	err := C.je_mallctl(
		dirtyKey,
		nil,
		nil,
		unsafe.Pointer(&cDirty),
		C.size_t(unsafe.Sizeof(cDirty)),
	)
	if err != 0 {
		return fmt.Errorf("failed to set dirty_decay_ms: error %d", err)
	}

	// Update arenas.muzzy_decay_ms
	muzzyKey := C.CString("arenas.muzzy_decay_ms")
	defer C.free(unsafe.Pointer(muzzyKey))

	err = C.je_mallctl(
		muzzyKey,
		nil,
		nil,
		unsafe.Pointer(&cMuzzy),
		C.size_t(unsafe.Sizeof(cMuzzy)),
	)
	if err != 0 {
		return fmt.Errorf("failed to set muzzy_decay_ms: error %d", err)
	}

	return nil
}

snippet-5.go

package jemalloc

/*
#cgo LDFLAGS: -ljemalloc
#include <stdlib.h>
#include <jemalloc/jemalloc.h>
*/
import "C"
import (
	"fmt"
	"unsafe"
)

// JemallocStats contains memory metrics collected from jemalloc.
type JemallocStats struct {
	Allocated uint64 // Total bytes allocated to the application
	Active    uint64 // Total bytes in active pages mapped by jemalloc
	Mapped    uint64 // Total bytes mapped by the allocator from the OS
}

// ReadStats forces a statistic refresh and returns the current jemalloc stats.
func ReadStats() (*JemallocStats, error) {
	// 1. Refresh jemalloc's cached statistics by updating the epoch
	epochKey := C.CString("epoch")
	defer C.free(unsafe.Pointer(epochKey))

	var epoch uint64 = 1
	epochSize := C.size_t(unsafe.Sizeof(epoch))
	C.je_mallctl(epochKey, nil, nil, unsafe.Pointer(&epoch), epochSize)

	var stats JemallocStats

	// 2. Query stats.allocated
	allocKey := C.CString("stats.allocated")
	defer C.free(unsafe.Pointer(allocKey))
	var allocated C.size_t
	size := C.size_t(unsafe.Sizeof(allocated))
	if err := C.je_mallctl(allocKey, unsafe.Pointer(&allocated), &size, nil, 0); err != 0 {
		return nil, fmt.Errorf("failed to query stats.allocated: %d", err)
	}
	stats.Allocated = uint64(allocated)

	// 3. Query stats.active
	activeKey := C.CString("stats.active")
	defer C.free(unsafe.Pointer(activeKey))
	var active C.size_t
	size = C.size_t(unsafe.Sizeof(active))
	if err := C.je_mallctl(activeKey, unsafe.Pointer(&active), &size, nil, 0); err != 0 {
		return nil, fmt.Errorf("failed to query stats.active: %d", err)
	}
	stats.Active = uint64(active)

	// 4. Query stats.mapped
	mappedKey := C.CString("stats.mapped")
	defer C.free(unsafe.Pointer(mappedKey))
	var mapped C.size_t
	size = C.size_t(unsafe.Sizeof(mapped))
	if err := C.je_mallctl(mappedKey, unsafe.Pointer(&mapped), &size, nil, 0); err != 0 {
		return nil, fmt.Errorf("failed to query stats.mapped: %d", err)
	}
	stats.Mapped = uint64(mapped)

	return &stats, nil
}

snippet-6.go

package jemalloc

/*
#cgo LDFLAGS: -ljemalloc
#include <stdlib.h>
#include <jemalloc/jemalloc.h>
*/
import "C"
import (
	"fmt"
	"unsafe"
)

// DumpProfile triggers an immediate write of the jemalloc heap profile to the target filepath.
// Profiling must be enabled at startup via MALLOC_CONF="prof:true".
func DumpProfile(filepath string) error {
	dumpKey := C.CString("prof.dump")
	defer C.free(unsafe.Pointer(dumpKey))

	cPath := C.CString(filepath)
	defer C.free(unsafe.Pointer(cPath))

	err := C.je_mallctl(
		dumpKey,
		nil,
		nil,
		unsafe.Pointer(&cPath),
		C.size_t(unsafe.Sizeof(cPath)),
	)
	if err != 0 {
		return fmt.Errorf("failed to dump jemalloc profile: %d", err)
	}

	return nil
}

snippet-8.go

package main

import (
	"fmt"
	"io"
	"net/http"
	"strconv"

	"jemalloc" // Replace with the import path of your wrapper package
)

func fileUploadHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	contentLengthStr := r.Header.Get("Content-Length")
	if contentLengthStr == "" {
		http.Error(w, "Missing Content-Length header", http.StatusLengthRequired)
		return
	}

	contentLength, err := strconv.Atoi(contentLengthStr)
	if err != nil || contentLength <= 0 {
		http.Error(w, "Invalid Content-Length", http.StatusBadRequest)
		return
	}

	// 1. Allocate an off-heap buffer via jemalloc
	buf := jemalloc.NewBuffer(contentLength)
	defer buf.Release()

	// 2. Read request body directly into the off-heap slice
	slice := buf.Bytes()
	_, err = io.ReadFull(r.Body, slice)
	if err != nil && err != io.EOF {
		http.Error(w, "Failed to read request body", http.StatusInternalServerError)
		return
	}

	// 3. Process the data (e.g., validate, parse JSON, or stream to disk)
	// Passing the slice off-heap ensures Go's GC never sweeps this memory.
	fmt.Printf("Successfully read %d bytes off-heap\n", len(slice))

	w.WriteHeader(http.StatusOK)
	w.Write([]byte("Upload processed successfully"))
}