132 lines
3.4 KiB
Go
132 lines
3.4 KiB
Go
package main
|
|
|
|
import (
|
|
"flag"
|
|
"log"
|
|
"net/http"
|
|
"os"
|
|
"os/signal"
|
|
"strconv"
|
|
"syscall"
|
|
"time"
|
|
|
|
"github.com/prometheus/client_golang/prometheus"
|
|
"github.com/prometheus/client_golang/prometheus/collectors"
|
|
"github.com/prometheus/client_golang/prometheus/promhttp"
|
|
)
|
|
|
|
var (
|
|
allocatedBytes = prometheus.NewGauge(prometheus.GaugeOpts{
|
|
Name: "memconsumer_allocated_bytes",
|
|
Help: "Total bytes allocated by the memory consumer",
|
|
})
|
|
allocationCount = prometheus.NewCounter(prometheus.CounterOpts{
|
|
Name: "memconsumer_allocations_total",
|
|
Help: "Total number of memory allocations performed",
|
|
})
|
|
)
|
|
|
|
// memoryHolder prevents GC from reclaiming allocations
|
|
var memoryHolder [][]byte
|
|
|
|
func envInt(key string, fallback int) int {
|
|
if v := os.Getenv(key); v != "" {
|
|
if i, err := strconv.Atoi(v); err == nil {
|
|
return i
|
|
}
|
|
}
|
|
return fallback
|
|
}
|
|
|
|
func envDuration(key string, fallback time.Duration) time.Duration {
|
|
if v := os.Getenv(key); v != "" {
|
|
if d, err := time.ParseDuration(v); err == nil {
|
|
return d
|
|
}
|
|
}
|
|
return fallback
|
|
}
|
|
|
|
func envString(key, fallback string) string {
|
|
if v := os.Getenv(key); v != "" {
|
|
return v
|
|
}
|
|
return fallback
|
|
}
|
|
|
|
func main() {
|
|
// Env vars set defaults, flags can override
|
|
initialMB := flag.Int("initial-mb", envInt("INITIAL_MB", 0), "Initial memory to allocate (MB) [env: INITIAL_MB]")
|
|
allocMB := flag.Int("alloc-mb", envInt("ALLOC_MB", 10), "Memory to allocate per iteration (MB) [env: ALLOC_MB]")
|
|
maxMB := flag.Int("max-mb", envInt("MAX_MB", 1000), "Max total memory to allocate, 0=unlimited (MB) [env: MAX_MB]")
|
|
interval := flag.Duration("interval", envDuration("INTERVAL", 5*time.Second), "Allocation interval [env: INTERVAL]")
|
|
metricsAddr := flag.String("metrics-addr", envString("METRICS_ADDR", ":8080"), "Prometheus metrics address [env: METRICS_ADDR]")
|
|
flag.Parse()
|
|
|
|
var totalAllocated int
|
|
|
|
reg := prometheus.NewRegistry()
|
|
reg.MustRegister(
|
|
collectors.NewGoCollector(),
|
|
collectors.NewProcessCollector(collectors.ProcessCollectorOpts{}),
|
|
allocatedBytes,
|
|
allocationCount,
|
|
)
|
|
|
|
// Initial allocation
|
|
if *initialMB > 0 {
|
|
allocate(*initialMB)
|
|
totalAllocated += *initialMB
|
|
log.Printf("Initial allocation: %d MB", *initialMB)
|
|
}
|
|
|
|
// Metrics server
|
|
http.Handle("/metrics", promhttp.HandlerFor(reg, promhttp.HandlerOpts{}))
|
|
http.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(http.StatusOK)
|
|
})
|
|
go func() {
|
|
log.Printf("Metrics server listening on %s", *metricsAddr)
|
|
if err := http.ListenAndServe(*metricsAddr, nil); err != nil {
|
|
log.Fatalf("metrics server failed: %v", err)
|
|
}
|
|
}()
|
|
|
|
// Graceful shutdown
|
|
stop := make(chan os.Signal, 1)
|
|
signal.Notify(stop, syscall.SIGTERM, syscall.SIGINT)
|
|
|
|
ticker := time.NewTicker(*interval)
|
|
defer ticker.Stop()
|
|
|
|
log.Printf("Starting: alloc=%dMB interval=%v max=%dMB", *allocMB, *interval, *maxMB)
|
|
|
|
for {
|
|
select {
|
|
case <-ticker.C:
|
|
if *maxMB > 0 && totalAllocated+*allocMB > *maxMB {
|
|
log.Printf("Max limit reached (%d MB), skipping allocation", *maxMB)
|
|
continue
|
|
}
|
|
allocate(*allocMB)
|
|
totalAllocated += *allocMB
|
|
log.Printf("Allocated %d MB, total held: %d MB", *allocMB, totalAllocated)
|
|
case <-stop:
|
|
log.Println("Shutting down")
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func allocate(mb int) {
|
|
size := mb * 1024 * 1024
|
|
buf := make([]byte, size)
|
|
// Touch pages to ensure RSS allocation
|
|
for i := 0; i < size; i += 4096 {
|
|
buf[i] = 1
|
|
}
|
|
memoryHolder = append(memoryHolder, buf)
|
|
allocatedBytes.Add(float64(size))
|
|
allocationCount.Inc()
|
|
}
|