Skip to content

Commit a9dff7f

Browse files
joshhaugjoshhaug
authored
Merge pull request #1630 from NASA-AMMOS/feature/streaming-thread
Stream simulation results in their own thread
2 parents 532fcda + 76d1831 commit a9dff7f

File tree

2 files changed

+200
-159
lines changed

2 files changed

+200
-159
lines changed

‎merlin-worker/src/main/java/gov/nasa/jpl/aerie/merlin/worker/postgres/PostgresProfileQueryHandler.java

+183
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,183 @@
1+
package gov.nasa.jpl.aerie.merlin.worker.postgres;
2+
3+
import gov.nasa.jpl.aerie.json.JsonParser;
4+
import gov.nasa.jpl.aerie.merlin.driver.resources.ResourceProfile;
5+
import gov.nasa.jpl.aerie.merlin.driver.resources.ResourceProfiles;
6+
import gov.nasa.jpl.aerie.merlin.protocol.types.Duration;
7+
import gov.nasa.jpl.aerie.merlin.protocol.types.RealDynamics;
8+
import gov.nasa.jpl.aerie.merlin.protocol.types.SerializedValue;
9+
import gov.nasa.jpl.aerie.merlin.server.remotes.postgres.DatabaseException;
10+
import gov.nasa.jpl.aerie.merlin.server.remotes.postgres.FailedInsertException;
11+
import gov.nasa.jpl.aerie.merlin.server.remotes.postgres.FailedUpdateException;
12+
import gov.nasa.jpl.aerie.merlin.server.remotes.postgres.PreparedStatements;
13+
import org.apache.commons.lang3.tuple.Pair;
14+
15+
import javax.sql.DataSource;
16+
import java.sql.Connection;
17+
import java.sql.PreparedStatement;
18+
import java.sql.SQLException;
19+
import java.sql.Statement;
20+
import java.util.HashMap;
21+
22+
import static gov.nasa.jpl.aerie.merlin.driver.json.SerializedValueJsonParser.serializedValueP;
23+
import static gov.nasa.jpl.aerie.merlin.server.http.ProfileParsers.realDynamicsP;
24+
import static gov.nasa.jpl.aerie.merlin.server.remotes.postgres.PostgresParsers.discreteProfileTypeP;
25+
import static gov.nasa.jpl.aerie.merlin.server.remotes.postgres.PostgresParsers.realProfileTypeP;
26+
27+
/**
28+
* Utility class to handle upload of resource profiles to the database.
29+
* */
30+
public class PostgresProfileQueryHandler implements AutoCloseable {
31+
private final Connection connection;
32+
private final HashMap<String, Integer> profileIds;
33+
private final HashMap<String, Duration> profileDurations;
34+
35+
private final PreparedStatement postProfileStatement;
36+
private final PreparedStatement postSegmentsStatement;
37+
private final PreparedStatement updateDurationStatement;
38+
39+
public PostgresProfileQueryHandler(DataSource dataSource, long datasetId) throws SQLException {
40+
connection = dataSource.getConnection();
41+
profileIds = new HashMap<>();
42+
profileDurations = new HashMap<>();
43+
44+
final String postProfilesSql =
45+
//language=sql
46+
"""
47+
insert into merlin.profile (dataset_id, name, type, duration)
48+
values (%d, ?, ?::jsonb, ?::interval)
49+
on conflict (dataset_id, name) do nothing
50+
""".formatted(datasetId);
51+
final String postSegmentsSql =
52+
//language=sql
53+
"""
54+
insert into merlin.profile_segment (dataset_id, profile_id, start_offset, dynamics, is_gap)
55+
values (%d, ?, ?::interval, ?::jsonb, false)
56+
""".formatted(datasetId);
57+
final String updateDurationSql =
58+
//language=SQL
59+
"""
60+
update merlin.profile
61+
set duration = ?::interval
62+
where (dataset_id, id) = (%d, ?);
63+
""".formatted(datasetId);
64+
65+
postProfileStatement = connection.prepareStatement(postProfilesSql, PreparedStatement.RETURN_GENERATED_KEYS);
66+
postSegmentsStatement = connection.prepareStatement(postSegmentsSql, PreparedStatement.NO_GENERATED_KEYS);
67+
updateDurationStatement = connection.prepareStatement(updateDurationSql, PreparedStatement.NO_GENERATED_KEYS);
68+
}
69+
70+
/**
71+
* Upload profiles, profile segments, and corresponding profile durations to the database.
72+
* */
73+
public void uploadResourceProfiles(final ResourceProfiles resourceProfiles) {
74+
try {
75+
// Add new profiles to DB
76+
for (final var realEntry : resourceProfiles.realProfiles().entrySet()) {
77+
if (!profileIds.containsKey(realEntry.getKey())) {
78+
addRealProfileToBatch(realEntry.getKey(), realEntry.getValue());
79+
}
80+
}
81+
for (final var discreteEntry : resourceProfiles.discreteProfiles().entrySet()) {
82+
if (!profileIds.containsKey(discreteEntry.getKey())) {
83+
addDiscreteProfileToBatch(discreteEntry.getKey(), discreteEntry.getValue());
84+
}
85+
}
86+
postProfiles();
87+
88+
// Post Segments
89+
for (final var realEntry : resourceProfiles.realProfiles().entrySet()) {
90+
addProfileSegmentsToBatch(realEntry.getKey(), realEntry.getValue(), realDynamicsP);
91+
}
92+
for (final var discreteEntry : resourceProfiles.discreteProfiles().entrySet()) {
93+
addProfileSegmentsToBatch(discreteEntry.getKey(), discreteEntry.getValue(), serializedValueP);
94+
}
95+
96+
postProfileSegments();
97+
updateProfileDurations();
98+
} catch (SQLException ex) {
99+
throw new DatabaseException("Exception occurred while posting profiles.", ex);
100+
}
101+
}
102+
103+
private void addRealProfileToBatch(final String name, ResourceProfile<RealDynamics> profile) throws SQLException {
104+
postProfileStatement.setString(1, name);
105+
postProfileStatement.setString(2, realProfileTypeP.unparse(Pair.of("real", profile.schema())).toString());
106+
PreparedStatements.setDuration(this.postProfileStatement, 3, Duration.ZERO);
107+
108+
postProfileStatement.addBatch();
109+
110+
profileDurations.put(name, Duration.ZERO);
111+
}
112+
113+
private void addDiscreteProfileToBatch(final String name, ResourceProfile<SerializedValue> profile) throws SQLException {
114+
postProfileStatement.setString(1, name);
115+
postProfileStatement.setString(2, discreteProfileTypeP.unparse(Pair.of("discrete", profile.schema())).toString());
116+
PreparedStatements.setDuration(this.postProfileStatement, 3, Duration.ZERO);
117+
118+
postProfileStatement.addBatch();
119+
120+
profileDurations.put(name, Duration.ZERO);
121+
}
122+
123+
/**
124+
* Insert the batched profiles and cache their ids for future use.
125+
*
126+
* This method takes advantage of the fact that we're using the Postgres JDBC,
127+
* which returns all columns when executing batches with `getGeneratedKeys`.
128+
*/
129+
private void postProfiles() throws SQLException {
130+
final var results = this.postProfileStatement.executeBatch();
131+
for (final var result : results) {
132+
if (result == Statement.EXECUTE_FAILED) throw new FailedInsertException("merlin.profile_segment");
133+
}
134+
135+
final var resultSet = this.postProfileStatement.getGeneratedKeys();
136+
while(resultSet.next()){
137+
profileIds.put(resultSet.getString("name"), resultSet.getInt("id"));
138+
}
139+
}
140+
141+
private void postProfileSegments() throws SQLException {
142+
final var results = this.postSegmentsStatement.executeBatch();
143+
for (final var result : results) {
144+
if (result == Statement.EXECUTE_FAILED) throw new FailedInsertException("merlin.profile_segment");
145+
}
146+
}
147+
148+
private void updateProfileDurations() throws SQLException {
149+
final var results = this.updateDurationStatement.executeBatch();
150+
for (final var result : results) {
151+
if (result == Statement.EXECUTE_FAILED) throw new FailedUpdateException("merlin.profile");
152+
}
153+
}
154+
155+
private <T> void addProfileSegmentsToBatch(final String name, ResourceProfile<T> profile, JsonParser<T> dynamicsP) throws SQLException {
156+
final var id = profileIds.get(name);
157+
this.postSegmentsStatement.setLong(1, id);
158+
159+
var newDuration = profileDurations.get(name);
160+
for (final var segment : profile.segments()) {
161+
PreparedStatements.setDuration(this.postSegmentsStatement, 2, newDuration);
162+
final var dynamics = dynamicsP.unparse(segment.dynamics()).toString();
163+
this.postSegmentsStatement.setString(3, dynamics);
164+
this.postSegmentsStatement.addBatch();
165+
166+
newDuration = newDuration.plus(segment.extent());
167+
}
168+
169+
this.updateDurationStatement.setLong(2, id);
170+
PreparedStatements.setDuration(this.updateDurationStatement, 1, newDuration);
171+
this.updateDurationStatement.addBatch();
172+
173+
profileDurations.put(name, newDuration);
174+
}
175+
176+
@Override
177+
public void close() throws SQLException {
178+
this.postProfileStatement.close();
179+
this.postSegmentsStatement.close();
180+
this.updateDurationStatement.close();
181+
this.connection.close();
182+
}
183+
}

‎merlin-worker/src/main/java/gov/nasa/jpl/aerie/merlin/worker/postgres/PostgresProfileStreamer.java

+17-159
Original file line numberDiff line numberDiff line change
@@ -1,180 +1,38 @@
11
package gov.nasa.jpl.aerie.merlin.worker.postgres;
22

3-
import gov.nasa.jpl.aerie.json.JsonParser;
4-
import gov.nasa.jpl.aerie.merlin.driver.resources.ResourceProfile;
53
import gov.nasa.jpl.aerie.merlin.driver.resources.ResourceProfiles;
6-
import gov.nasa.jpl.aerie.merlin.protocol.types.Duration;
7-
import gov.nasa.jpl.aerie.merlin.protocol.types.RealDynamics;
8-
import gov.nasa.jpl.aerie.merlin.protocol.types.SerializedValue;
94
import gov.nasa.jpl.aerie.merlin.server.remotes.postgres.DatabaseException;
10-
import gov.nasa.jpl.aerie.merlin.server.remotes.postgres.FailedInsertException;
11-
import gov.nasa.jpl.aerie.merlin.server.remotes.postgres.FailedUpdateException;
12-
import gov.nasa.jpl.aerie.merlin.server.remotes.postgres.PreparedStatements;
13-
import org.apache.commons.lang3.tuple.Pair;
145

156
import javax.sql.DataSource;
16-
import java.sql.Connection;
17-
import java.sql.PreparedStatement;
187
import java.sql.SQLException;
19-
import java.sql.Statement;
20-
import java.util.HashMap;
8+
import java.util.concurrent.ExecutorService;
9+
import java.util.concurrent.Executors;
2110
import java.util.function.Consumer;
2211

23-
import static gov.nasa.jpl.aerie.merlin.server.remotes.postgres.PostgresParsers.discreteProfileTypeP;
24-
import static gov.nasa.jpl.aerie.merlin.server.remotes.postgres.PostgresParsers.realProfileTypeP;
25-
26-
import static gov.nasa.jpl.aerie.merlin.driver.json.SerializedValueJsonParser.serializedValueP;
27-
import static gov.nasa.jpl.aerie.merlin.server.http.ProfileParsers.realDynamicsP;
28-
2912
public class PostgresProfileStreamer implements Consumer<ResourceProfiles>, AutoCloseable {
30-
private final Connection connection;
31-
private final HashMap<String, Integer> profileIds;
32-
private final HashMap<String, Duration> profileDurations;
33-
34-
private final PreparedStatement postProfileStatement;
35-
private final PreparedStatement postSegmentsStatement;
36-
private final PreparedStatement updateDurationStatement;
13+
private final ExecutorService queryQueue;
14+
private final PostgresProfileQueryHandler queryHandler;
3715

3816
public PostgresProfileStreamer(DataSource dataSource, long datasetId) throws SQLException {
39-
this.connection = dataSource.getConnection();
40-
profileIds = new HashMap<>();
41-
profileDurations = new HashMap<>();
42-
43-
final String postProfilesSql =
44-
//language=sql
45-
"""
46-
insert into merlin.profile (dataset_id, name, type, duration)
47-
values (%d, ?, ?::jsonb, ?::interval)
48-
on conflict (dataset_id, name) do nothing
49-
""".formatted(datasetId);
50-
final String postSegmentsSql =
51-
//language=sql
52-
"""
53-
insert into merlin.profile_segment (dataset_id, profile_id, start_offset, dynamics, is_gap)
54-
values (%d, ?, ?::interval, ?::jsonb, false)
55-
""".formatted(datasetId);
56-
final String updateDurationSql =
57-
//language=SQL
58-
"""
59-
update merlin.profile
60-
set duration = ?::interval
61-
where (dataset_id, id) = (%d, ?);
62-
""".formatted(datasetId);
63-
64-
postProfileStatement = connection.prepareStatement(postProfilesSql, PreparedStatement.RETURN_GENERATED_KEYS);
65-
postSegmentsStatement = connection.prepareStatement(postSegmentsSql, PreparedStatement.NO_GENERATED_KEYS);
66-
updateDurationStatement = connection.prepareStatement(updateDurationSql, PreparedStatement.NO_GENERATED_KEYS);
17+
this.queryQueue = Executors.newSingleThreadExecutor();
18+
this.queryHandler = new PostgresProfileQueryHandler(dataSource, datasetId);
6719
}
6820

6921
@Override
7022
public void accept(final ResourceProfiles resourceProfiles) {
71-
try {
72-
// Add new profiles to DB
73-
for(final var realEntry : resourceProfiles.realProfiles().entrySet()){
74-
if(!profileIds.containsKey(realEntry.getKey())){
75-
addRealProfileToBatch(realEntry.getKey(), realEntry.getValue());
76-
}
77-
}
78-
for(final var discreteEntry : resourceProfiles.discreteProfiles().entrySet()) {
79-
if(!profileIds.containsKey(discreteEntry.getKey())){
80-
addDiscreteProfileToBatch(discreteEntry.getKey(), discreteEntry.getValue());
81-
}
82-
}
83-
postProfiles();
84-
85-
// Post Segments
86-
for(final var realEntry : resourceProfiles.realProfiles().entrySet()){
87-
addProfileSegmentsToBatch(realEntry.getKey(), realEntry.getValue(), realDynamicsP);
88-
}
89-
for(final var discreteEntry : resourceProfiles.discreteProfiles().entrySet()) {
90-
addProfileSegmentsToBatch(discreteEntry.getKey(), discreteEntry.getValue(), serializedValueP);
91-
}
92-
93-
postProfileSegments();
94-
updateProfileDurations();
95-
} catch (SQLException ex) {
96-
throw new DatabaseException("Exception occurred while posting profiles.", ex);
97-
}
98-
}
99-
100-
private void addRealProfileToBatch(final String name, ResourceProfile<RealDynamics> profile) throws SQLException {
101-
postProfileStatement.setString(1, name);
102-
postProfileStatement.setString(2, realProfileTypeP.unparse(Pair.of("real", profile.schema())).toString());
103-
PreparedStatements.setDuration(this.postProfileStatement, 3, Duration.ZERO);
104-
105-
postProfileStatement.addBatch();
106-
107-
profileDurations.put(name, Duration.ZERO);
108-
}
109-
110-
private void addDiscreteProfileToBatch(final String name, ResourceProfile<SerializedValue> profile) throws SQLException {
111-
postProfileStatement.setString(1, name);
112-
postProfileStatement.setString(2, discreteProfileTypeP.unparse(Pair.of("discrete", profile.schema())).toString());
113-
PreparedStatements.setDuration(this.postProfileStatement, 3, Duration.ZERO);
114-
115-
postProfileStatement.addBatch();
116-
117-
profileDurations.put(name, Duration.ZERO);
118-
}
119-
120-
/**
121-
* Insert the batched profiles and cache their ids for future use.
122-
*
123-
* This method takes advantage of the fact that we're using the Postgres JDBC,
124-
* which returns all columns when executing batches with `getGeneratedKeys`.
125-
*/
126-
private void postProfiles() throws SQLException {
127-
final var results = this.postProfileStatement.executeBatch();
128-
for (final var result : results) {
129-
if (result == Statement.EXECUTE_FAILED) throw new FailedInsertException("merlin.profile_segment");
130-
}
131-
132-
final var resultSet = this.postProfileStatement.getGeneratedKeys();
133-
while(resultSet.next()){
134-
profileIds.put(resultSet.getString("name"), resultSet.getInt("id"));
135-
}
136-
}
137-
138-
private void postProfileSegments() throws SQLException {
139-
final var results = this.postSegmentsStatement.executeBatch();
140-
for (final var result : results) {
141-
if (result == Statement.EXECUTE_FAILED) throw new FailedInsertException("merlin.profile_segment");
142-
}
143-
}
144-
145-
private void updateProfileDurations() throws SQLException {
146-
final var results = this.updateDurationStatement.executeBatch();
147-
for (final var result : results) {
148-
if (result == Statement.EXECUTE_FAILED) throw new FailedUpdateException("merlin.profile");
149-
}
150-
}
151-
152-
private <T> void addProfileSegmentsToBatch(final String name, ResourceProfile<T> profile, JsonParser<T> dynamicsP) throws SQLException {
153-
final var id = profileIds.get(name);
154-
this.postSegmentsStatement.setLong(1, id);
155-
156-
var newDuration = profileDurations.get(name);
157-
for (final var segment : profile.segments()) {
158-
PreparedStatements.setDuration(this.postSegmentsStatement, 2, newDuration);
159-
final var dynamics = dynamicsP.unparse(segment.dynamics()).toString();
160-
this.postSegmentsStatement.setString(3, dynamics);
161-
this.postSegmentsStatement.addBatch();
162-
163-
newDuration = newDuration.plus(segment.extent());
164-
}
165-
166-
this.updateDurationStatement.setLong(2, id);
167-
PreparedStatements.setDuration(this.updateDurationStatement, 1, newDuration);
168-
this.updateDurationStatement.addBatch();
169-
170-
profileDurations.put(name, newDuration);
23+
queryQueue.submit(() -> {
24+
queryHandler.uploadResourceProfiles(resourceProfiles);
25+
});
17126
}
17227

17328
@Override
174-
public void close() throws SQLException {
175-
this.postProfileStatement.close();
176-
this.postSegmentsStatement.close();
177-
this.updateDurationStatement.close();
178-
this.connection.close();
29+
public void close() {
30+
queryQueue.shutdown();
31+
try {
32+
queryHandler.close();
33+
} catch (SQLException e) {
34+
throw new DatabaseException("Error occurred while attempting to close PostgresProfileQueryHandler", e);
35+
}
17936
}
37+
18038
}

0 commit comments

Comments
 (0)