poc: a metrics module for pebble

internals/daemon/api.go

@@ -108,6 +108,10 @@ var API = []*Command{{
 	WriteAccess: AdminAccess{},
 	GET:         v1GetIdentities,
 	POST:        v1PostIdentities,
+}, {
+	Path:       "/metrics",
+	ReadAccess: OpenAccess{},
+	GET:        Metrics,
 }}
 
 var (

internals/daemon/api_metrics.go

@@ -0,0 +1,35 @@
+// Copyright (c) 2024 Canonical Ltd
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License version 3 as
+// published by the Free Software Foundation.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+package daemon
+
+import (
+	"net/http"
+
+	"github.com/canonical/pebble/internals/metrics"
+)
+
+func Metrics(c *Command, r *http.Request, _ *UserState) Response {
+	return metricsResponse{}
+}
+
+// metricsResponse is a Response implementation to serve the metrics in a prometheus metrics format.
+type metricsResponse struct{}
+
+func (r metricsResponse) ServeHTTP(w http.ResponseWriter, req *http.Request) {
+	registry := metrics.GetRegistry()
+	w.WriteHeader(http.StatusOK)
+	w.Write([]byte(registry.GatherMetrics()))
+
+}

internals/daemon/daemon.go

@@ -20,6 +20,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"math/rand"
 	"net"
 	"net/http"
 	"os"
@@ -36,6 +37,7 @@ import (
 	"gopkg.in/tomb.v2"
 
 	"github.com/canonical/pebble/internals/logger"
+	"github.com/canonical/pebble/internals/metrics"
 	"github.com/canonical/pebble/internals/osutil"
 	"github.com/canonical/pebble/internals/osutil/sys"
 	"github.com/canonical/pebble/internals/overlord"
@@ -366,6 +368,20 @@ func (d *Daemon) Init() error {
 	}
 
 	logger.Noticef("Started daemon.")
+
+	registry := metrics.GetRegistry()
+	registry.NewMetric("my_counter", metrics.MetricTypeCounter, "A simple counter")
+	// Goroutine to update metrics randomly
+	go func() {
+		for {
+			err := registry.IncCounter("my_counter") // Increment by 1
+			if err != nil {
+				fmt.Println("Error incrementing counter:", err)
+			}
+			time.Sleep(time.Duration(rand.Intn(5)+1) * time.Second) // Random sleep between 1 and 5 seconds
+		}
+	}()
+
 	return nil
 }
 

internals/metrics/metrics.go

@@ -0,0 +1,211 @@
+// Copyright (c) 2024 Canonical Ltd
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License version 3 as
+// published by the Free Software Foundation.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+package metrics
+
+import (
+	"fmt"
+	"math/rand"
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/gorilla/mux"
+)
+
+// MetricType models the type of a metric.
+type MetricType string
+
+const (
+	MetricTypeCounter   MetricType = "counter"
+	MetricTypeGauge     MetricType = "gauge"
+	MetricTypeHistogram MetricType = "histogram"
+)
+
+// Metric represents a single metric.
+type Metric struct {
+	Name  string
+	Type  MetricType
+	Help  string
+	value interface{} // Can be int64 for counter/gauge, or []float64 for histogram.
+	mu    sync.RWMutex
+}
+
+// MetricsRegistry stores and manages metrics.
+type MetricsRegistry struct {
+	metrics map[string]*Metric
+	mu      sync.RWMutex
+}
+
+// Package-level variable to hold the single registry.
+var registry *MetricsRegistry
+
+// Ensure registry is initialized only once.
+var once sync.Once
+
+// GetRegistry returns the singleton MetricsRegistry instance.
+func GetRegistry() *MetricsRegistry {
+	once.Do(func() {
+		registry = &MetricsRegistry{
+			metrics: make(map[string]*Metric),
+		}
+	})
+	return registry
+}
+
+// NewMetric registers a new metric.
+func (r *MetricsRegistry) NewMetric(name string, metricType MetricType, help string) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	if _, ok := r.metrics[name]; ok {
+		return fmt.Errorf("metric with name %s already registered", name)
+	}
+
+	var value interface{}
+	switch metricType {
+	case MetricTypeCounter, MetricTypeGauge:
+		value = int64(0)
+	case MetricTypeHistogram:
+		value = make([]float64, 0)
+	default:
+		return fmt.Errorf("invalid metric type: %s", metricType)
+	}
+
+	r.metrics[name] = &Metric{
+		Name:  name,
+		Type:  metricType,
+		Help:  help,
+		value: value,
+	}
+
+	return nil
+}
+
+// IncCounter increments a counter metric.
+func (r *MetricsRegistry) IncCounter(name string) error {
+	return r.updateMetric(name, MetricTypeCounter, 1)
+}
+
+// SetGauge sets the value of a gauge metric.
+func (r *MetricsRegistry) SetGauge(name string, value int64) error {
+	return r.updateMetric(name, MetricTypeGauge, value)
+}
+
+// ObserveHistogram adds a value to a histogram metric.
+func (r *MetricsRegistry) ObserveHistogram(name string, value float64) error {
+	return r.updateMetric(name, MetricTypeHistogram, value)
+}
+
+// updateMetric updates the value of a metric.
+func (r *MetricsRegistry) updateMetric(name string, metricType MetricType, value interface{}) error {
+	r.mu.RLock()
+	metric, ok := r.metrics[name]
+	r.mu.RUnlock()
+
+	if !ok {
+		return fmt.Errorf("metric with name %s not found", name)
+	}
+
+	if metric.Type != metricType {
+		return fmt.Errorf("mismatched metric type for %s", name)
+	}
+
+	metric.mu.Lock()
+	defer metric.mu.Unlock()
+
+	switch metricType {
+	case MetricTypeCounter:
+		metric.value = metric.value.(int64) + int64(value.(int))
+	case MetricTypeGauge:
+		metric.value = value.(int64)
+	case MetricTypeHistogram:
+		metric.value = append(metric.value.([]float64), value.(float64))
+	}
+
+	return nil
+}
+
+// GatherMetrics gathers all metrics and formats them in Prometheus exposition format.
+func (r *MetricsRegistry) GatherMetrics() string {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+
+	var output string
+	for _, metric := range r.metrics {
+		output += fmt.Sprintf("# HELP %s %s\n", metric.Name, metric.Help)
+		output += fmt.Sprintf("# TYPE %s %s\n", metric.Name, metric.Type)
+
+		switch metric.Type {
+		case MetricTypeCounter, MetricTypeGauge:
+			output += fmt.Sprintf("%s %d\n", metric.Name, metric.value.(int64))
+		case MetricTypeHistogram:
+			for _, v := range metric.value.([]float64) {
+				output += fmt.Sprintf("%s %f\n", metric.Name, v) // Basic histogram representation
+			}
+		}
+	}
+
+	return output
+}
+
+// Example usage
+func main() {
+	// Get the singleton registry
+	registry := GetRegistry()
+
+	registry.NewMetric("my_counter", MetricTypeCounter, "A simple counter")
+	registry.NewMetric("my_gauge", MetricTypeGauge, "A simple gauge")
+	registry.NewMetric("my_histogram", MetricTypeHistogram, "A simple histogram")
+
+	// Goroutine to update metrics randomly
+	go func() {
+		for {
+			// Counter
+			err := registry.IncCounter("my_counter") // Increment by 1
+			if err != nil {
+				fmt.Println("Error incrementing counter:", err)
+			}
+
+			// Gauge
+			gaugeValue := rand.Int63n(100)
+			err = registry.SetGauge("my_gauge", gaugeValue)
+			if err != nil {
+				fmt.Println("Error setting gauge:", err)
+			}
+
+			// Histogram
+			histogramValue := rand.Float64() * 10
+			err = registry.ObserveHistogram("my_histogram", histogramValue)
+			if err != nil {
+				fmt.Println("Error observing histogram:", err)
+			}
+
+			time.Sleep(time.Duration(rand.Intn(5)+1) * time.Second) // Random sleep between 1 and 5 seconds
+		}
+	}()
+	// Use Gorilla Mux router
+	router := mux.NewRouter()
+	router.HandleFunc("/metrics", func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte(registry.GatherMetrics()))
+	})
+
+	// Serve on port 2112
+	fmt.Println("Serving metrics on :2112/metrics")
+	err := http.ListenAndServe(":2112", router) // Use the router here
+	if err != nil {
+		fmt.Println("Error starting server:", err)
+	}
+}

internals/metrics/metrics_test.go

@@ -0,0 +1,85 @@
+// Copyright (c) 2024 Canonical Ltd
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License version 3 as
+// published by the Free Software Foundation.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+package metrics
+
+import (
+	"sync"
+	"testing"
+
+	. "gopkg.in/check.v1"
+)
+
+// Hook up check.v1 into the "go test" runner
+func Test(t *testing.T) { TestingT(t) }
+
+var _ = Suite(&RegistryTestSuite{})
+
+// Test Suite structure
+type RegistryTestSuite struct {
+	registry *MetricsRegistry
+}
+
+func (s *RegistryTestSuite) SetUpTest(c *C) {
+	s.registry = &MetricsRegistry{
+		metrics: make(map[string]*Metric),
+	}
+}
+
+func (s *RegistryTestSuite) TestCounter(c *C) {
+	s.registry.NewMetric("test_counter", MetricTypeCounter, "Test counter")
+	s.registry.IncCounter("test_counter")
+	s.registry.IncCounter("test_counter")
+	c.Check(s.registry.metrics["test_counter"].value.(int64), Equals, int64(2))
+}
+
+func (s *RegistryTestSuite) TestGauge(c *C) {
+	s.registry.NewMetric("test_gauge", MetricTypeGauge, "Test gauge")
+	s.registry.SetGauge("test_gauge", 10)
+	c.Check(s.registry.metrics["test_gauge"].value.(int64), Equals, int64(10))
+	s.registry.SetGauge("test_gauge", 20)
+	c.Check(s.registry.metrics["test_gauge"].value.(int64), Equals, int64(20))
+}
+
+func (s *RegistryTestSuite) TestHistogram(c *C) {
+	s.registry.NewMetric("test_histogram", MetricTypeHistogram, "Test histogram")
+	s.registry.ObserveHistogram("test_histogram", 1.0)
+	s.registry.ObserveHistogram("test_histogram", 2.0)
+	histogramValues := s.registry.metrics["test_histogram"].value.([]float64)
+	c.Check(len(histogramValues), Equals, 2)
+	c.Check(histogramValues[0], Equals, 1.0)
+	c.Check(histogramValues[1], Equals, 2.0)
+}
+
+func (s *RegistryTestSuite) TestGatherMetrics(c *C) {
+	s.registry.NewMetric("test_counter", MetricTypeCounter, "Test counter")
+	s.registry.IncCounter("test_counter")
+	metricsOutput := s.registry.GatherMetrics()
+	expectedOutput := "# HELP test_counter Test counter\n# TYPE test_counter counter\ntest_counter 1\n"
+	c.Check(metricsOutput, Equals, expectedOutput)
+}
+
+func (s *RegistryTestSuite) TestRaceConditions(c *C) {
+	s.registry.NewMetric("race_counter", MetricTypeCounter, "Race counter")
+	var wg sync.WaitGroup
+	for i := 0; i < 1000; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			s.registry.IncCounter("race_counter")
+		}()
+	}
+	wg.Wait()
+	c.Check(s.registry.metrics["race_counter"].value.(int64), Equals, int64(1000))
+}