10
10
import java .io .IOException ;
11
11
import java .io .OutputStream ;
12
12
import java .net .InetSocketAddress ;
13
+ import java .nio .charset .StandardCharsets ;
13
14
import java .util .ArrayList ;
14
15
import java .util .List ;
15
- import java .util .concurrent .Executors ;
16
+ import java .util .concurrent .* ;
16
17
import java .util .function .BooleanSupplier ;
17
18
18
19
public class HealthServer {
@@ -21,7 +22,9 @@ public class HealthServer {
21
22
public final int port ;
22
23
public final String endpoint ;
23
24
public final HttpServer httpServer ;
24
- private List <BooleanSupplier > checks = new ArrayList <>();
25
+ private final ExecutorService healthCheckExecutor =
26
+ Executors .newCachedThreadPool ();
27
+ private final List <BooleanSupplier > checks = new CopyOnWriteArrayList <>();
25
28
26
29
public HealthServer (final int port , @ NotNull final String endpoint ) throws IOException {
27
30
this .port = port ;
@@ -30,14 +33,14 @@ public HealthServer(final int port, @NotNull final String endpoint) throws IOExc
30
33
httpServer = HttpServer .create (new InetSocketAddress (port ), 0 );
31
34
httpServer .createContext ("/" , createDefaultHandler ());
32
35
httpServer .createContext (endpoint , createHandler ());
33
- httpServer .setExecutor (Executors . newSingleThreadExecutor () );
36
+ httpServer .setExecutor (healthCheckExecutor );
34
37
httpServer .start ();
35
38
log .info ("HealthServer started" );
36
39
}
37
40
38
41
private void writeResponse (@ NotNull final HttpExchange httpExchange , @ NotNull final int responseCode , @ NotNull final String responseBody ) throws IOException {
39
- final byte [] response = responseBody .getBytes ("UTF-8" );
40
- httpExchange .getResponseHeaders ().add ("Content-Type" , "text/plain; charset=UTF-8" );
42
+ final byte [] response = responseBody .getBytes (StandardCharsets . UTF_8 );
43
+ httpExchange .getResponseHeaders ().add ("Content-Type" , "text/plain; charset=" + StandardCharsets . UTF_8 . name () );
41
44
httpExchange .sendResponseHeaders (responseCode , response .length );
42
45
final OutputStream out = httpExchange .getResponseBody ();
43
46
out .write (response );
@@ -91,16 +94,67 @@ public void clearChecks() {
91
94
}
92
95
93
96
public boolean checkHealth () {
94
- boolean isHealthy = true ;
95
- for (final BooleanSupplier check : checks ) {
96
- isHealthy &= check .getAsBoolean ();
97
+ try {
98
+ CompletionService <Boolean > executorCompletionService
99
+ = new ExecutorCompletionService <>(healthCheckExecutor );
100
+ int n = checks .size ();
101
+ List <Future <Boolean >> futures = new ArrayList <>(n );
102
+ try {
103
+ for (BooleanSupplier check : checks ) {
104
+ futures .add (executorCompletionService .submit (checkToCallable (check )));
105
+ }
106
+ for (int i = 0 ; i < n ; ++i ) {
107
+ try {
108
+ Boolean result = executorCompletionService .take ().get ();
109
+ if (result == null || !result ) {
110
+ return false ; // Return false immediately if any check fails
111
+ }
112
+ } catch (ExecutionException e ) {
113
+ log .error ("A health check task execution failed. Marking unhealthy." , e .getCause () != null ? e .getCause () : e );
114
+ return false ;
115
+ } catch (InterruptedException e ) {
116
+ log .error ("Health check interrupted. Marking unhealthy." , e );
117
+ Thread .currentThread ().interrupt ();
118
+ return false ;
119
+ }
120
+ }
121
+ } finally {
122
+ for (Future <Boolean > f : futures ) {
123
+ f .cancel (true );
124
+ }
125
+ }
126
+ return true ; // Return true only if all checks pass
127
+ } catch (Exception e ) {
128
+ log .error ("Exception during health checks" , e );
129
+ return false ;
97
130
}
98
- return isHealthy ;
99
131
}
100
132
101
133
public void close () {
102
134
if (httpServer != null ) {
103
135
httpServer .stop (0 );
104
136
}
137
+ if (healthCheckExecutor != null ) {
138
+ healthCheckExecutor .shutdown ();
139
+ try {
140
+ if (!healthCheckExecutor .awaitTermination (3 , TimeUnit .SECONDS )) {
141
+ healthCheckExecutor .shutdownNow ();
142
+ }
143
+ } catch (InterruptedException ie ) {
144
+ healthCheckExecutor .shutdownNow ();
145
+ Thread .currentThread ().interrupt ();
146
+ }
147
+ }
148
+ }
149
+
150
+ private static Callable <Boolean > checkToCallable (BooleanSupplier check ) {
151
+ return () -> {
152
+ try {
153
+ return check .getAsBoolean ();
154
+ } catch (Exception e ) {
155
+ log .error ("Exception during health check" , e );
156
+ return false ;
157
+ }
158
+ };
105
159
}
106
160
}
0 commit comments