initial commit, but the application might be already complete ;-)

2026-01-08 15:31:28 +01:00
commit f068a6b4e8
7 changed files with 307 additions and 0 deletions
--- a/main.go
+++ b/main.go
@@ -0,0 +1,131 @@
+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()
+}