memory.go raw

   1  //go:build !(js && wasm)
   2  
   3  package ratelimit
   4  
   5  import (
   6  	"errors"
   7  	"runtime"
   8  
   9  	"github.com/pbnjay/memory"
  10  )
  11  
  12  // MinimumMemoryMB is the minimum memory required to run the relay with rate limiting.
  13  const MinimumMemoryMB = 128
  14  
  15  // AutoDetectMemoryFraction is the fraction of available memory to use when auto-detecting.
  16  const AutoDetectMemoryFraction = 0.66
  17  
  18  // DefaultMaxMemoryMB is the default maximum memory target when auto-detecting.
  19  // This caps the auto-detected value to ensure optimal performance.
  20  const DefaultMaxMemoryMB = 1500
  21  
  22  // ErrInsufficientMemory is returned when there isn't enough memory to run the relay.
  23  var ErrInsufficientMemory = errors.New("insufficient memory: relay requires at least 128MB of available memory")
  24  
  25  // ProcessMemoryStats contains memory statistics for the current process.
  26  // On Linux, these are read from /proc/self/status for accurate RSS values.
  27  // On other platforms, these are approximated from Go runtime stats.
  28  type ProcessMemoryStats struct {
  29  	// VmRSS is the resident set size (total physical memory in use) in bytes
  30  	VmRSS uint64
  31  	// RssShmem is the shared memory portion of RSS in bytes
  32  	RssShmem uint64
  33  	// RssAnon is the anonymous (non-shared) memory in bytes
  34  	RssAnon uint64
  35  	// VmHWM is the peak RSS (high water mark) in bytes
  36  	VmHWM uint64
  37  }
  38  
  39  // PhysicalMemoryBytes returns the actual physical memory usage (RSS - shared)
  40  func (p ProcessMemoryStats) PhysicalMemoryBytes() uint64 {
  41  	if p.VmRSS > p.RssShmem {
  42  		return p.VmRSS - p.RssShmem
  43  	}
  44  	return p.VmRSS
  45  }
  46  
  47  // PhysicalMemoryMB returns the actual physical memory usage in MB
  48  func (p ProcessMemoryStats) PhysicalMemoryMB() uint64 {
  49  	return p.PhysicalMemoryBytes() / (1024 * 1024)
  50  }
  51  
  52  // DetectAvailableMemoryMB returns the available system memory in megabytes.
  53  // On Linux, this returns the actual available memory (free + cached).
  54  // On other systems, it returns total memory minus the Go runtime's current usage.
  55  func DetectAvailableMemoryMB() uint64 {
  56  	// Use pbnjay/memory for cross-platform memory detection
  57  	available := memory.FreeMemory()
  58  	if available == 0 {
  59  		// Fallback: use total memory
  60  		available = memory.TotalMemory()
  61  	}
  62  	return available / (1024 * 1024)
  63  }
  64  
  65  // DetectTotalMemoryMB returns the total system memory in megabytes.
  66  func DetectTotalMemoryMB() uint64 {
  67  	return memory.TotalMemory() / (1024 * 1024)
  68  }
  69  
  70  // CalculateTargetMemoryMB calculates the target memory limit based on configuration.
  71  // If configuredMB is 0, it auto-detects based on available memory (66% of available, capped at 1.5GB).
  72  // If configuredMB is non-zero, it validates that it's achievable.
  73  // Returns an error if there isn't enough memory.
  74  func CalculateTargetMemoryMB(configuredMB int) (int, error) {
  75  	availableMB := int(DetectAvailableMemoryMB())
  76  
  77  	// If configured to auto-detect (0), calculate target
  78  	if configuredMB == 0 {
  79  		// First check if we have minimum available memory
  80  		if availableMB < MinimumMemoryMB {
  81  			return 0, ErrInsufficientMemory
  82  		}
  83  
  84  		// Calculate 66% of available
  85  		targetMB := int(float64(availableMB) * AutoDetectMemoryFraction)
  86  
  87  		// If 66% is less than minimum, use minimum (we've already verified we have enough)
  88  		if targetMB < MinimumMemoryMB {
  89  			targetMB = MinimumMemoryMB
  90  		}
  91  
  92  		// Cap at default maximum for optimal performance
  93  		if targetMB > DefaultMaxMemoryMB {
  94  			targetMB = DefaultMaxMemoryMB
  95  		}
  96  
  97  		return targetMB, nil
  98  	}
  99  
 100  	// If explicitly configured, validate it's achievable
 101  	if configuredMB < MinimumMemoryMB {
 102  		return 0, ErrInsufficientMemory
 103  	}
 104  
 105  	// Warn but allow if configured target exceeds available
 106  	// (the PID controller will throttle as needed)
 107  	return configuredMB, nil
 108  }
 109  
 110  // GetMemoryStats returns current memory statistics for logging.
 111  type MemoryStats struct {
 112  	TotalMB       uint64
 113  	AvailableMB   uint64
 114  	TargetMB      int
 115  	GoAllocatedMB uint64
 116  	GoSysMB       uint64
 117  }
 118  
 119  // GetMemoryStats returns current memory statistics.
 120  func GetMemoryStats(targetMB int) MemoryStats {
 121  	var m runtime.MemStats
 122  	runtime.ReadMemStats(&m)
 123  
 124  	return MemoryStats{
 125  		TotalMB:       DetectTotalMemoryMB(),
 126  		AvailableMB:   DetectAvailableMemoryMB(),
 127  		TargetMB:      targetMB,
 128  		GoAllocatedMB: m.Alloc / (1024 * 1024),
 129  		GoSysMB:       m.Sys / (1024 * 1024),
 130  	}
 131  }
 132  
 133  // readProcessMemoryStatsFallback returns memory stats using Go runtime.
 134  // This is used on non-Linux platforms or when /proc is unavailable.
 135  // The values are approximations and may not accurately reflect OS-level metrics.
 136  func readProcessMemoryStatsFallback() ProcessMemoryStats {
 137  	var m runtime.MemStats
 138  	runtime.ReadMemStats(&m)
 139  
 140  	// Use Sys as an approximation of RSS (includes all memory from OS)
 141  	// HeapAlloc approximates anonymous memory (live heap objects)
 142  	// We cannot determine shared memory from Go runtime, so leave it at 0
 143  	return ProcessMemoryStats{
 144  		VmRSS:    m.Sys,
 145  		RssAnon:  m.HeapAlloc,
 146  		RssShmem: 0, // Cannot determine shared memory from Go runtime
 147  		VmHWM:    0, // Not available from Go runtime
 148  	}
 149  }
 150