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