From 0f99c431981ef741241f122b46e587dec3da7a2a Mon Sep 17 00:00:00 2001
From: Mattia Cansirro Cortorillo <mcansirro@doxee.com>
Date: Wed, 18 May 2022 20:50:15 +0200
Subject: [PATCH] Add PUT shutdown operation for Prometheus Push Gateway

See gh-31104
---
 .../PrometheusPushGatewayManager.java         | 19 ++++++++++++++++++-
 .../PrometheusPushGatewayManagerTests.java    | 10 ++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/metrics/export/prometheus/PrometheusPushGatewayManager.java b/spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/metrics/export/prometheus/PrometheusPushGatewayManager.java
index e02182783e..b0e2540f88 100644
--- a/spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/metrics/export/prometheus/PrometheusPushGatewayManager.java
+++ b/spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/metrics/export/prometheus/PrometheusPushGatewayManager.java
@@ -110,6 +110,15 @@ public class PrometheusPushGatewayManager {
 		}
 	}
 
+	private void put() {
+		try {
+			this.pushGateway.push(this.registry, this.job, this.groupingKey);
+		}
+		catch (Throwable ex) {
+			logger.warn("Unexpected exception thrown while pushing metrics to Prometheus Pushgateway", ex);
+		}
+	}
+
 	private void delete() {
 		try {
 			this.pushGateway.delete(this.job, this.groupingKey);
@@ -135,6 +144,9 @@ public class PrometheusPushGatewayManager {
 		case PUSH:
 			push();
 			break;
+		case PUT:
+			put();
+			break;
 		case DELETE:
 			delete();
 			break;
@@ -152,10 +164,15 @@ public class PrometheusPushGatewayManager {
 		NONE,
 
 		/**
-		 * Perform a 'push' before shutdown.
+		 * Perform a 'push' with POST method before shutdown.
 		 */
 		PUSH,
 
+		/**
+		 * Perform a 'push' with PUT method before shutdown.
+		 */
+		PUT,
+
 		/**
 		 * Perform a 'delete' before shutdown.
 		 */
diff --git a/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/metrics/export/prometheus/PrometheusPushGatewayManagerTests.java b/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/metrics/export/prometheus/PrometheusPushGatewayManagerTests.java
index 7bb3ceef2f..f06d633fab 100644
--- a/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/metrics/export/prometheus/PrometheusPushGatewayManagerTests.java
+++ b/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/metrics/export/prometheus/PrometheusPushGatewayManagerTests.java
@@ -144,6 +144,16 @@ class PrometheusPushGatewayManagerTests {
 		then(this.pushGateway).should().pushAdd(this.registry, "job", this.groupingKey);
 	}
 
+	@Test
+	void shutdownWhenShutdownOperationIsPutPerformsPushOnShutdown() throws Exception {
+		givenScheduleAtFixedRateWithReturnFuture();
+		PrometheusPushGatewayManager manager = new PrometheusPushGatewayManager(this.pushGateway, this.registry,
+				this.scheduler, this.pushRate, "job", this.groupingKey, ShutdownOperation.PUT);
+		manager.shutdown();
+		then(this.future).should().cancel(false);
+		then(this.pushGateway).should().push(this.registry, "job", this.groupingKey);
+	}
+
 	@Test
 	void shutdownWhenShutdownOperationIsDeletePerformsDeleteOnShutdown() throws Exception {
 		givenScheduleAtFixedRateWithReturnFuture();