Source code

001/*
002 * Copyright © 2025 CUI-OpenSource-Software (info@cuioss.de)
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 *     http://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016package de.cuioss.benchmarking.common.metrics;
017
018import de.cuioss.benchmarking.common.config.BenchmarkConfiguration;
019import de.cuioss.tools.logging.CuiLogger;
020import org.openjdk.jmh.results.RunResult;
021
022import java.io.IOException;
023import java.nio.file.Files;
024import java.nio.file.Path;
025import java.time.Duration;
026import java.time.Instant;
027import java.util.Collection;
028import java.util.List;
029import java.util.Map;
030import java.util.concurrent.ConcurrentHashMap;
031
032import static de.cuioss.benchmarking.common.util.BenchmarkingLogMessages.ERROR.FAILED_COLLECT_PROMETHEUS_BENCHMARK;
033import static de.cuioss.benchmarking.common.util.BenchmarkingLogMessages.ERROR.PROMETHEUS_CONNECTIVITY_FAILED;
034import static de.cuioss.benchmarking.common.util.BenchmarkingLogMessages.INFO.*;
035import static de.cuioss.benchmarking.common.util.BenchmarkingLogMessages.WARN.*;
036
037/**
038 * Centralized manager for Prometheus metrics collection during benchmark execution.
039 * This class provides a unified approach to collecting real-time metrics from Prometheus
040 * for both JMH and WRK benchmarks.
041 *
042 * <p>The manager tracks benchmark execution timestamps and collects metrics
043 * at the end of benchmark runs. This ensures consistent metrics collection
044 * across different benchmark types.</p>
045 *
046 * <p>Features:</p>
047 * <ul>
048 *   <li>Thread-safe timestamp tracking for concurrent benchmarks</li>
049 *   <li>Automatic Prometheus client configuration from system properties</li>
050 *   <li>Unified metrics collection for different benchmark result types</li>
051 *   <li>Integration with MetricsOrchestrator for actual metrics retrieval</li>
052 * </ul>
053 */
054public class PrometheusMetricsManager {
055
056    private static final CuiLogger LOGGER = new CuiLogger(PrometheusMetricsManager.class);
057
058    public static final String PROMETHEUS_DIR_NAME = "prometheus";
059    public static final String PROMETHEUS_URL_PROPERTY = "prometheus.url";
060    public static final String PROMETHEUS_URL_DEFAULT = "http://localhost:9090";
061    public static final String METRICS_FILE_SUFFIX = "-metrics.json";
062
063    private final Map<String, BenchmarkTimestamps> timestampTracker = new ConcurrentHashMap<>();
064    private final String prometheusUrl;
065    private final boolean metricsEnabled;
066
067    /**
068     * Data class to hold benchmark execution timestamps.
069     */
070    private record BenchmarkTimestamps(Instant startTime, Instant endTime) {
071        boolean isValid() {
072            return startTime != null && endTime != null;
073        }
074    }
075
076    /**
077     * Creates a new PrometheusMetricsManager.
078     * Prometheus URL is read from system property or defaults to localhost:9090.
079     */
080    public PrometheusMetricsManager() {
081        this.prometheusUrl = System.getProperty(PROMETHEUS_URL_PROPERTY, PROMETHEUS_URL_DEFAULT);
082        this.metricsEnabled = isPrometheusAvailable();
083
084        if (metricsEnabled) {
085            LOGGER.info(PROMETHEUS_ENABLED, prometheusUrl);
086        } else {
087            LOGGER.info(PROMETHEUS_DISABLED, prometheusUrl);
088        }
089    }
090
091    /**
092     * Records the start time for a benchmark execution.
093     * This should be called just before benchmark execution begins.
094     *
095     * @param benchmarkName the name of the benchmark
096     */
097    public void recordBenchmarkStart(String benchmarkName) {
098        Instant startTime = Instant.now();
099        timestampTracker.compute(benchmarkName, (key, existing) -> {
100            if (existing != null) {
101                LOGGER.warn(OVERWRITING_START_TIME, benchmarkName);
102            }
103            return new BenchmarkTimestamps(startTime, null);
104        });
105        LOGGER.debug("Benchmark '%s' started at: %s", benchmarkName, startTime);
106    }
107
108    /**
109     * Records both start and end timestamps for a benchmark.
110     * Used when timestamps are provided externally (e.g., from WRK output).
111     *
112     * @param benchmarkName the name of the benchmark
113     * @param startTime the benchmark start time
114     * @param endTime the benchmark end time
115     */
116    public void recordBenchmarkTimestamps(String benchmarkName, Instant startTime, Instant endTime) {
117        if (startTime == null || endTime == null) {
118            LOGGER.warn(INVALID_TIMESTAMPS, benchmarkName, startTime, endTime);
119            return;
120        }
121
122        timestampTracker.put(benchmarkName, new BenchmarkTimestamps(startTime, endTime));
123        LOGGER.debug("Recorded timestamps for benchmark '%s': %s to %s",
124                benchmarkName, startTime, endTime);
125    }
126
127    /**
128     * Collects Prometheus metrics for completed benchmarks.
129     * This should be called after all benchmarks have completed execution.
130     *
131     * @param results the benchmark results
132     * @param config the benchmark configuration
133     */
134    public void collectMetricsForResults(Collection<RunResult> results, BenchmarkConfiguration config) {
135        if (!metricsEnabled) {
136            LOGGER.debug("Skipping Prometheus metrics collection - not enabled");
137            return;
138        }
139
140        if (!config.hasIntegrationConfig()) {
141            LOGGER.debug("Skipping Prometheus metrics collection - no integration config");
142            return;
143        }
144
145        String outputDirectory = config.resultsDirectory();
146
147        try {
148            Path prometheusDir = Path.of(outputDirectory, PROMETHEUS_DIR_NAME);
149            Files.createDirectories(prometheusDir);
150
151            LOGGER.info(COLLECTING_PROMETHEUS_METRICS, prometheusUrl);
152
153            MetricsOrchestrator orchestrator = new MetricsOrchestrator(
154                    new PrometheusClient(prometheusUrl)
155            );
156
157            for (RunResult result : results) {
158                String benchmarkName = extractBenchmarkName(result);
159                collectMetricsForBenchmark(benchmarkName, orchestrator, prometheusDir);
160            }
161
162        } catch (IOException e) {
163            LOGGER.warn(FAILED_COLLECT_PROMETHEUS, e.getMessage());
164        }
165    }
166
167    /**
168     * Collects Prometheus metrics for WRK benchmark results.
169     * This method handles the specific case of WRK benchmarks where
170     * timestamps are extracted from the output files.
171     *
172     * @param benchmarkName the name of the benchmark
173     * @param startTime the benchmark start time
174     * @param endTime the benchmark end time
175     * @param outputDirectory the directory to save metrics
176     */
177    public void collectMetricsForWrkBenchmark(String benchmarkName, Instant startTime,
178            Instant endTime, String outputDirectory) {
179        if (!metricsEnabled) {
180            LOGGER.warn(SKIPPING_PROMETHEUS_COLLECTION, prometheusUrl);
181            LOGGER.debug("To enable metrics collection, ensure Prometheus is running and accessible at the configured URL");
182            return;
183        }
184
185        if (startTime == null || endTime == null) {
186            LOGGER.warn(MISSING_TIMESTAMPS_FOR_COLLECTION, benchmarkName);
187            return;
188        }
189
190        try {
191            Path prometheusDir = Path.of(outputDirectory, PROMETHEUS_DIR_NAME);
192            Files.createDirectories(prometheusDir);
193
194            LOGGER.info(COLLECTING_WRK_PROMETHEUS_METRICS, benchmarkName);
195
196            MetricsOrchestrator orchestrator = new MetricsOrchestrator(
197                    new PrometheusClient(prometheusUrl)
198            );
199
200            orchestrator.collectBenchmarkMetrics(
201                    benchmarkName,
202                    startTime,
203                    endTime,
204                    prometheusDir
205            );
206
207            LOGGER.info(PROMETHEUS_METRICS_SAVED, prometheusDir, benchmarkName, METRICS_FILE_SUFFIX);
208
209        } catch (IOException e) {
210            LOGGER.error(e, FAILED_COLLECT_PROMETHEUS_BENCHMARK, benchmarkName, prometheusUrl);
211            LOGGER.debug("Attempted to query metrics for time range: %s to %s", startTime, endTime);
212        }
213    }
214
215    private void collectMetricsForBenchmark(String benchmarkName, MetricsOrchestrator orchestrator,
216            Path prometheusDir) {
217        LOGGER.debug("Collecting metrics for benchmark: '%s', available keys: %s",
218                benchmarkName, timestampTracker.keySet());
219
220        // Get the specific timestamps for this benchmark
221        BenchmarkTimestamps timestamps = timestampTracker.get(benchmarkName);
222
223        if (timestamps == null || !timestamps.isValid()) {
224            LOGGER.warn(NO_VALID_TIMESTAMPS, benchmarkName, timestampTracker.keySet());
225            return;
226        }
227
228        LOGGER.info(USING_SESSION_TIMESTAMPS, benchmarkName, timestamps.startTime(), timestamps.endTime());
229
230        try {
231            LOGGER.info(COLLECTING_BENCHMARK_METRICS, benchmarkName);
232
233            orchestrator.collectBenchmarkMetrics(
234                    benchmarkName,
235                    timestamps.startTime(),
236                    timestamps.endTime(),
237                    prometheusDir
238            );
239
240            LOGGER.info(PROMETHEUS_METRICS_SAVED, prometheusDir, benchmarkName, METRICS_FILE_SUFFIX);
241
242        } catch (IOException e) {
243            LOGGER.warn(FAILED_COLLECT_BENCHMARK_METRICS, benchmarkName, e.getMessage());
244        }
245    }
246
247    private String extractBenchmarkName(RunResult result) {
248        String label = result.getPrimaryResult().getLabel();
249        int lastDot = label.lastIndexOf('.');
250        if (lastDot >= 0 && lastDot < label.length() - 1) {
251            return label.substring(lastDot + 1);
252        }
253        return label;
254    }
255
256    private boolean isPrometheusAvailable() {
257        try {
258            LOGGER.debug("Checking Prometheus availability at URL: %s", prometheusUrl);
259            PrometheusClient client = new PrometheusClient(prometheusUrl);
260            client.queryRange(List.of("up"), Instant.now().minusSeconds(60), Instant.now(), Duration.ofSeconds(10));
261            LOGGER.debug("Prometheus is available and responding at: %s", prometheusUrl);
262            return true;
263        } catch (PrometheusClient.PrometheusException e) {
264            LOGGER.error(e, PROMETHEUS_CONNECTIVITY_FAILED, prometheusUrl, e.getMessage(), e.getClass().getSimpleName());
265            LOGGER.info(PROMETHEUS_CONNECTIVITY_ADVICE, prometheusUrl);
266            return false;
267        }
268    }
269
270    /**
271     * Clears all tracked timestamps.
272     * Useful for testing or when reusing the manager for multiple benchmark runs.
273     */
274    public void clear() {
275        timestampTracker.clear();
276    }
277
278}