Skip to content

Commit 1108545

Browse files
ybezsonovybezsonov
authored
Merge pull request #722 from aws-samples/bedrock-tda
Bedrock tda
2 parents 768ec55 + 788c9c9 commit 1108545

26 files changed

Lines changed: 2621 additions & 169 deletions

‎apps/unicorn-store-spring/pom.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,11 @@
5858
<groupId>org.springframework.boot</groupId>
5959
<artifactId>spring-boot-starter-actuator</artifactId>
6060
</dependency>
61+
<!-- https://mvnrepository.com/artifact/io.micrometer/micrometer-registry-prometheus -->
62+
<dependency>
63+
<groupId>io.micrometer</groupId>
64+
<artifactId>micrometer-registry-prometheus</artifactId>
65+
</dependency>
6166
<!-- AWS SDK dependencies -->
6267
<dependency>
6368
<groupId>software.amazon.awssdk</groupId>

‎apps/unicorn-store-spring/src/main/java/com/unicorn/store/config/MonitoringConfig.java‎

Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,168 @@
1+
package com.unicorn.store.config;
2+
3+
import com.fasterxml.jackson.databind.JsonNode;
4+
import com.fasterxml.jackson.databind.ObjectMapper;
5+
import io.micrometer.core.instrument.MeterRegistry;
6+
import io.micrometer.core.instrument.config.MeterFilter;
7+
import org.springframework.boot.actuate.autoconfigure.metrics.MeterRegistryCustomizer;
8+
import org.springframework.context.annotation.Bean;
9+
import org.springframework.context.annotation.Configuration;
10+
11+
import java.io.File;
12+
import java.io.IOException;
13+
import java.net.InetAddress;
14+
import java.net.URI;
15+
import java.net.UnknownHostException;
16+
import java.net.http.HttpClient;
17+
import java.net.http.HttpRequest;
18+
import java.net.http.HttpResponse;
19+
import java.nio.file.Files;
20+
import java.util.Optional;
21+
22+
@Configuration
23+
public class MonitoringConfig {
24+
25+
private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
26+
private static final File NAMESPACE_FILE = new File("/var/run/secrets/kubernetes.io/serviceaccount/namespace");
27+
28+
@Bean
29+
public MeterRegistryCustomizer<MeterRegistry> meterRegistryCustomizer() {
30+
return registry -> {
31+
String clusterType = System.getenv("ECS_CONTAINER_METADATA_URI_V4") != null ? "ecs" : "eks";
32+
String cluster = clusterType.equals("ecs") ? extractClusterNameFromMetadata().orElse("unknown") : Optional.ofNullable(System.getenv("CLUSTER")).orElse("unknown");
33+
String containerName = "unicorn-store-spring";
34+
String taskOrPodId = extractTaskOrPodId().orElse("unknown");
35+
String namespace = clusterType.equals("eks") ? readNamespaceFile().orElse("default") : "";
36+
37+
// Get the container/pod IP address
38+
String ipAddress = getContainerOrPodIp().orElse("unknown");
39+
40+
registry.config().commonTags(
41+
"cluster", cluster,
42+
"cluster_type", clusterType,
43+
"container_name", containerName,
44+
"task_pod_id", taskOrPodId,
45+
"instance", ipAddress, // Keep this for backward compatibility
46+
"container_ip", ipAddress // Add this new tag that won't be overwritten
47+
);
48+
49+
if (!namespace.isEmpty()) {
50+
registry.config().commonTags("namespace", namespace);
51+
} else {
52+
registry.config().commonTags("namespace", "<no namespace>");
53+
}
54+
55+
registry.config().meterFilter(
56+
MeterFilter.deny(id ->
57+
id.getName().equals("jvm.gc.pause") &&
58+
!id.getTags().stream().allMatch(tag ->
59+
tag.getKey().equals("action") ||
60+
tag.getKey().equals("cause") ||
61+
tag.getKey().equals("gc")
62+
)
63+
)
64+
);
65+
};
66+
}
67+
68+
private Optional<String> extractTaskOrPodId() {
69+
String metadataUri = System.getenv("ECS_CONTAINER_METADATA_URI_V4");
70+
if (metadataUri != null) {
71+
try {
72+
HttpRequest request = HttpRequest.newBuilder()
73+
.uri(URI.create(metadataUri + "/task"))
74+
.build();
75+
76+
HttpResponse<String> response = HttpClient.newHttpClient()
77+
.send(request, HttpResponse.BodyHandlers.ofString());
78+
79+
JsonNode root = OBJECT_MAPPER.readTree(response.body());
80+
String taskArn = root.path("TaskARN").asText();
81+
String[] parts = taskArn.split("/");
82+
return parts.length > 1 ? Optional.of(parts[parts.length - 1]) : Optional.empty();
83+
84+
} catch (IOException | InterruptedException e) {
85+
return Optional.empty();
86+
}
87+
}
88+
89+
// EKS fallback: read pod name from Downward API
90+
return readFile("/etc/podinfo/name");
91+
}
92+
93+
private Optional<String> extractClusterNameFromMetadata() {
94+
String metadataUri = System.getenv("ECS_CONTAINER_METADATA_URI_V4");
95+
if (metadataUri != null) {
96+
try {
97+
HttpRequest request = HttpRequest.newBuilder()
98+
.uri(URI.create(metadataUri + "/task"))
99+
.build();
100+
101+
HttpResponse<String> response = HttpClient.newHttpClient()
102+
.send(request, HttpResponse.BodyHandlers.ofString());
103+
104+
JsonNode root = OBJECT_MAPPER.readTree(response.body());
105+
String clusterArn = root.path("Cluster").asText();
106+
String[] parts = clusterArn.split("/");
107+
return parts.length > 1 ? Optional.of(parts[parts.length - 1]) : Optional.empty();
108+
109+
} catch (IOException | InterruptedException e) {
110+
return Optional.empty();
111+
}
112+
}
113+
return Optional.empty();
114+
}
115+
116+
private Optional<String> readNamespaceFile() {
117+
return readFile(NAMESPACE_FILE.getAbsolutePath());
118+
}
119+
120+
private Optional<String> readFile(String path) {
121+
try {
122+
return Optional.of(Files.readString(new File(path).toPath()).trim());
123+
} catch (IOException e) {
124+
return Optional.empty();
125+
}
126+
}
127+
128+
// New method to get the container or pod IP address
129+
private Optional<String> getContainerOrPodIp() {
130+
// For ECS
131+
String metadataUri = System.getenv("ECS_CONTAINER_METADATA_URI_V4");
132+
if (metadataUri != null) {
133+
try {
134+
HttpRequest request = HttpRequest.newBuilder()
135+
.uri(URI.create(metadataUri))
136+
.build();
137+
138+
HttpResponse<String> response = HttpClient.newHttpClient()
139+
.send(request, HttpResponse.BodyHandlers.ofString());
140+
141+
JsonNode root = OBJECT_MAPPER.readTree(response.body());
142+
143+
if (root.has("Networks") && root.path("Networks").isArray() && !root.path("Networks").isEmpty()) {
144+
JsonNode network = root.path("Networks").get(0);
145+
if (network.has("IPv4Addresses") && network.path("IPv4Addresses").isArray() &&
146+
!network.path("IPv4Addresses").isEmpty()) {
147+
return Optional.of(network.path("IPv4Addresses").get(0).asText());
148+
}
149+
}
150+
} catch (IOException | InterruptedException e) {
151+
// Fall through to next method
152+
}
153+
}
154+
155+
// For Kubernetes/EKS
156+
String podIp = System.getenv("KUBERNETES_POD_IP");
157+
if (podIp != null && !podIp.isEmpty()) {
158+
return Optional.of(podIp);
159+
}
160+
161+
// Try to get local IP as fallback
162+
try {
163+
return Optional.of(InetAddress.getLocalHost().getHostAddress());
164+
} catch (UnknownHostException e) {
165+
return Optional.empty();
166+
}
167+
}
168+
}

‎apps/unicorn-store-spring/src/main/java/com/unicorn/store/controller/ThreadManagementController.java‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
package com.unicorn.store.controller;
2+
3+
import com.unicorn.store.service.ThreadGeneratorService;
4+
import org.springframework.http.ResponseEntity;
5+
import org.springframework.web.bind.annotation.*;
6+
7+
@RestController
8+
@RequestMapping("/api/threads")
9+
public class ThreadManagementController {
10+
11+
private final ThreadGeneratorService threadGeneratorService;
12+
13+
public ThreadManagementController(ThreadGeneratorService threadGeneratorService) {
14+
this.threadGeneratorService = threadGeneratorService;
15+
}
16+
17+
@PostMapping("/start")
18+
public ResponseEntity<String> startThreads(@RequestParam(defaultValue = "500") int count) {
19+
try {
20+
threadGeneratorService.startThreads(count);
21+
return ResponseEntity.ok("Successfully started " + count + " threads");
22+
} catch (IllegalStateException e) {
23+
return ResponseEntity.badRequest().body(e.getMessage());
24+
}
25+
}
26+
27+
@PostMapping("/stop")
28+
public ResponseEntity<String> stopThreads() {
29+
try {
30+
threadGeneratorService.stopThreads();
31+
return ResponseEntity.ok("Successfully stopped all threads");
32+
} catch (IllegalStateException e) {
33+
return ResponseEntity.badRequest().body(e.getMessage());
34+
}
35+
}
36+
37+
@GetMapping("/count")
38+
public ResponseEntity<Integer> getThreadCount() {
39+
return ResponseEntity.ok(threadGeneratorService.getActiveThreadCount());
40+
}
41+
}

‎apps/unicorn-store-spring/src/main/java/com/unicorn/store/monitoring/ThreadMonitoringMBean.java‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
package com.unicorn.store.monitoring;
2+
3+
import com.unicorn.store.service.ThreadGeneratorService;
4+
import org.springframework.jmx.export.annotation.ManagedAttribute;
5+
import org.springframework.jmx.export.annotation.ManagedResource;
6+
import org.springframework.stereotype.Component;
7+
8+
@Component
9+
@ManagedResource(objectName = "com.unicorn.store:type=ThreadMonitoring,name=ThreadStats")
10+
public class ThreadMonitoringMBean {
11+
12+
private final ThreadGeneratorService threadGeneratorService;
13+
14+
public ThreadMonitoringMBean(ThreadGeneratorService threadGeneratorService) {
15+
this.threadGeneratorService = threadGeneratorService;
16+
}
17+
18+
@ManagedAttribute(description = "Number of active custom threads")
19+
public int getActiveThreadCount() {
20+
return threadGeneratorService.getActiveThreadCount();
21+
}
22+
}

‎apps/unicorn-store-spring/src/main/java/com/unicorn/store/service/ThreadGeneratorService.java‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
package com.unicorn.store.service;
2+
3+
import org.slf4j.Logger;
4+
import org.slf4j.LoggerFactory;
5+
import org.springframework.stereotype.Service;
6+
7+
import java.util.ArrayList;
8+
import java.util.List;
9+
import java.util.concurrent.atomic.AtomicBoolean;
10+
11+
@Service
12+
public class ThreadGeneratorService {
13+
private static final Logger logger = LoggerFactory.getLogger(ThreadGeneratorService.class);
14+
private final List<Thread> activeThreads = new ArrayList<>();
15+
private final AtomicBoolean running = new AtomicBoolean(false);
16+
17+
public synchronized void startThreads(int threadCount) {
18+
if (running.get()) {
19+
throw new IllegalStateException("Threads are already running");
20+
}
21+
22+
running.set(true);
23+
logger.info("Starting {} threads", threadCount);
24+
25+
for (int i = 0; i < threadCount; i++) {
26+
Thread thread = new Thread(new DummyWorkload(running));
27+
thread.setName("DummyThread-" + i);
28+
activeThreads.add(thread);
29+
thread.start();
30+
}
31+
32+
logger.info("Started {} threads", threadCount);
33+
}
34+
35+
public synchronized void stopThreads() {
36+
if (!running.get()) {
37+
throw new IllegalStateException("No threads are running");
38+
}
39+
40+
logger.info("Stopping {} threads", activeThreads.size());
41+
running.set(false);
42+
43+
// Wait for all threads to complete
44+
activeThreads.forEach(thread -> {
45+
try {
46+
thread.join(5000); // Wait up to 5 seconds for each thread
47+
} catch (InterruptedException e) {
48+
logger.warn("Interrupted while waiting for thread {} to stop", thread.getName());
49+
}
50+
});
51+
52+
activeThreads.clear();
53+
logger.info("All threads stopped");
54+
}
55+
56+
public int getActiveThreadCount() {
57+
return activeThreads.size();
58+
}
59+
60+
private static class DummyWorkload implements Runnable {
61+
private final AtomicBoolean running;
62+
63+
public DummyWorkload(AtomicBoolean running) {
64+
this.running = running;
65+
}
66+
67+
@Override
68+
public void run() {
69+
while (running.get()) {
70+
// Simulate some work
71+
try {
72+
// Calculate some dummy values to keep CPU busy
73+
double result = 0;
74+
for (int i = 0; i < 1000; i++) {
75+
result += Math.sqrt(i) * Math.random();
76+
}
77+
Thread.sleep(100); // Sleep to prevent excessive CPU usage
78+
} catch (InterruptedException e) {
79+
Thread.currentThread().interrupt();
80+
break;
81+
}
82+
}
83+
}
84+
}
85+
}

‎apps/unicorn-store-spring/src/main/resources/application.properties‎

Lines changed: 17 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,27 @@
11
spring.sql.init.mode=always
2-
# spring.datasource.url=jdbc:postgresql://localhost:5432/unicorns
32
spring.datasource.username=postgres
4-
# spring.datasource.password=postgres
53
spring.jpa.database-platform=org.hibernate.dialect.PostgreSQLDialect
64
server.error.include-message=always
75
spring.jpa.open-in-view=false
8-
spring.datasource.hikari.maximumPoolSize=1
9-
spring.datasource.hikari.data-source-properties.preparedStatementCacheQueries=0
10-
management.endpoint.health.probes.enabled=true
116

12-
# Disable Hibernate usage of JDBC metadata
137
spring.jpa.properties.hibernate.temp.use_jdbc_metadata_defaults=false
14-
15-
# Database initialization will typically be performed outside of Spring lifecycle
168
spring.jpa.hibernate.ddl-auto=none
9+
10+
spring.datasource.hikari.initialization-fail-timeout=0
11+
spring.datasource.hikari.maximumPoolSize=1
1712
spring.datasource.hikari.allow-pool-suspension=true
13+
spring.datasource.hikari.data-source-properties.preparedStatementCacheQueries=0
14+
15+
# Virtual Threads
16+
spring.threads.virtual.enabled=true
17+
18+
# Actuator config
19+
spring.jmx.enabled=true
20+
management.endpoints.web.exposure.include=threaddump,prometheus,health,info
21+
management.endpoint.health.probes.enabled=true
22+
management.endpoint.prometheus.access=unrestricted
23+
management.endpoint.health.group.liveness.include=livenessState
24+
management.endpoint.health.group.readiness.include=readinessState
1825

19-
# Enable virtual threads
20-
spring.threads.virtual.enabled=true
26+
management.metrics.enable.all=true
27+
management.metrics.tags.application=otel-jmx-unicorn-store

0 commit comments

Comments
 (0)