api.go

// Copyright (c) 2025 Nathan Winemiller
// Copyright (c) 2019 Dropbox, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// THIS FILE HAS BEEN MODIFIED from its original version.
// Changes: TODO(nwinemiller) - List changes here

package udprobe

import (
	"fmt"
	"net/http"
	"sync"

	"github.com/prometheus/client_golang/prometheus/promhttp"
)

// API represnts the HTTP server answering queries for collected data.
type API struct {
	summarizer *Summarizer
	server     *http.Server
	ts         TagSet
	handler    *http.ServeMux
	mutex      sync.RWMutex
}

// PromHandler handles requests for Prometheus metrics.
func (api *API) PromHandler() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// 1. Update the Prometheus metrics based on the summary.

		// Lock the existing summaries cache
		api.summarizer.CMutex.RLock()
		summaries := api.summarizer.Cache
		LogInfo(fmt.Sprintf("Found %d data points", len(summaries)))
		// Convert the summaries to Prometheus metrics
		api.mutex.RLock()
		p := &PrometheusMetricSetter{}
		EmitMetricsFromSummaries(summaries, api.ts, p)
		api.mutex.RUnlock()
		// Unlock the cache
		api.summarizer.CMutex.RUnlock()

		// 2. Delegate the request to the official promhttp.Handler()
		// This serve the actual Prometheus formatted output
		promhttp.Handler().ServeHTTP(w, r)
	})
}

// StatusHandler acts as a back healthcheck and simply returns 200 OK.
func (api *API) StatusHandler(rw http.ResponseWriter, request *http.Request) {
	fmt.Fprintf(rw, "ok")
}

// Stop will close down the server and cause Run to exit.
func (api *API) Stop() {
	err := api.server.Close()
	if err != nil {
		HandleMinorErrorMsg(err, "Error stopping API")
	}
	LogInfo("API Stopped")
}

// Run calls RunForever in a separate goroutine for non-blocking behavior.
func (api *API) Run() {
	// This basically just exists to be consistent with the existing pattern
	// while also allowing it to be run blocking if desired.
	go api.RunForever()
}

// MergeUpdateTagSet combines a provided TagSet with the existing one
func (api *API) MergeUpdateTagSet(t TagSet) {
	api.mutex.Lock()
	if api.ts == nil {
		api.ts = make(TagSet)
	}
	// Copy new entries into the existing TagSet
	// Allowing retention of existing entries, updating where needed, and adding new
	for k, v := range t {
		api.ts[k] = v
	}
	api.mutex.Unlock()
}

// RunForever sets up the handlers above and then listens for requests until
// stopped or a fatal error occurs.
//
// Calling this will block until stopped/crashed.
func (api *API) RunForever() {
	// Setup the handlers
	// TODO(nwinemiller): It might be better to move this elsewhere?
	api.setupHandlers()
	err := api.server.ListenAndServe()
	if err != nil && err != http.ErrServerClosed {
		HandleFatalErrorMsg(err, "API server failed")
	}
}

// SetupHandlers attaches the handlers above to the http server mux.
func (api *API) setupHandlers() {
	api.handler.HandleFunc("/status", api.StatusHandler)
	api.handler.Handle("/metrics", api.PromHandler())
}

// New returns an initialized API struct.
func NewAPI(s *Summarizer, t TagSet, addr string) *API {
	// TODO(nwinemiller): In the future, make these options that can be provided.
	handler := http.NewServeMux()
	server := &http.Server{
		Addr:    addr,
		Handler: handler,
	}
	if t == nil {
		t = make(TagSet)
	}
	RegisterPrometheus() // Register the necessary variables with the Prometheus handler.
	return &API{summarizer: s, ts: t, handler: handler, server: server}
}