Skip to content

Commit

Permalink
Initial implementation of an experimental queue optimisation using Po…
Browse files Browse the repository at this point in the history
…stgres LISTEN/NOTIFY
  • Loading branch information
bjpirt committed Feb 24, 2024
1 parent 06bc5fd commit 0bf5a3d
Show file tree
Hide file tree
Showing 7 changed files with 537 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -87,8 +87,9 @@ public PostgresExecutionDAO postgresExecutionDAO(
@DependsOn({"flywayForPrimaryDb"})
public PostgresQueueDAO postgresQueueDAO(
@Qualifier("postgresRetryTemplate") RetryTemplate retryTemplate,
ObjectMapper objectMapper) {
return new PostgresQueueDAO(retryTemplate, objectMapper, dataSource);
ObjectMapper objectMapper,
PostgresProperties properties) {
return new PostgresQueueDAO(retryTemplate, objectMapper, dataSource, properties);
}

@Bean
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,10 @@ public class PostgresProperties {

private Integer deadlockRetryMax = 3;

private boolean experimentalQueueNotify = false;

private Integer experimentalQueueNotifyStalePeriod = 5000;

public String schema = "public";

public Duration getTaskDefCacheRefreshInterval() {
Expand All @@ -52,4 +56,20 @@ public String getSchema() {
public void setSchema(String schema) {
this.schema = schema;
}

public boolean getExperimentalQueueNotify() {
return experimentalQueueNotify;
}

public void setExperimentalQueueNotify(boolean experimentalQueueNotify) {
this.experimentalQueueNotify = experimentalQueueNotify;
}

public Integer getExperimentalQueueNotifyStalePeriod() {
return experimentalQueueNotifyStalePeriod;
}

public void setExperimentalQueueNotifyStalePeriod(Integer experimentalQueueNotifyStalePeriod) {
this.experimentalQueueNotifyStalePeriod = experimentalQueueNotifyStalePeriod;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,9 @@

import com.netflix.conductor.core.events.queue.Message;
import com.netflix.conductor.dao.QueueDAO;
import com.netflix.conductor.postgres.config.PostgresProperties;
import com.netflix.conductor.postgres.util.ExecutorsUtil;
import com.netflix.conductor.postgres.util.PostgresQueueListener;
import com.netflix.conductor.postgres.util.Query;

import com.fasterxml.jackson.databind.ObjectMapper;
Expand All @@ -40,8 +42,13 @@ public class PostgresQueueDAO extends PostgresBaseDAO implements QueueDAO {

private final ScheduledExecutorService scheduledExecutorService;

private PostgresQueueListener queueListener;

public PostgresQueueDAO(
RetryTemplate retryTemplate, ObjectMapper objectMapper, DataSource dataSource) {
RetryTemplate retryTemplate,
ObjectMapper objectMapper,
DataSource dataSource,
PostgresProperties properties) {
super(retryTemplate, objectMapper, dataSource);

this.scheduledExecutorService =
Expand All @@ -53,6 +60,10 @@ public PostgresQueueDAO(
UNACK_SCHEDULE_MS,
TimeUnit.MILLISECONDS);
logger.debug("{} is ready to serve", PostgresQueueDAO.class.getName());

if (properties.getExperimentalQueueNotify()) {
this.queueListener = new PostgresQueueListener(dataSource, properties);
}
}

@PreDestroy
Expand Down Expand Up @@ -169,6 +180,13 @@ public void remove(String queueName, String messageId) {

@Override
public int getSize(String queueName) {
if (queueListener != null) {
Integer size = queueListener.getSize(queueName);
if (size != null) {
return size;
}
}

final String GET_QUEUE_SIZE = "SELECT COUNT(*) FROM queue_message WHERE queue_name = ?";
return queryWithTransaction(
GET_QUEUE_SIZE, q -> ((Long) q.addParameter(queueName).executeCount()).intValue());
Expand Down Expand Up @@ -425,6 +443,12 @@ private boolean removeMessage(Connection connection, String queueName, String me
private List<Message> popMessages(
Connection connection, String queueName, int count, int timeout) {

if (this.queueListener != null) {
if (!this.queueListener.hasMessagesReady(queueName)) {
return new ArrayList<>();
}
}

String POP_QUERY =
"UPDATE queue_message SET popped = true WHERE message_id IN ("
+ "SELECT message_id FROM queue_message WHERE queue_name = ? AND popped = false AND "
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,207 @@
/*
* Copyright 2024 Conductor Authors.
* <p>
* Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
* <p>
* http://www.apache.org/licenses/LICENSE-2.0
* <p>
* Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
* an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
* specific language governing permissions and limitations under the License.
*/
package com.netflix.conductor.postgres.util;

import java.sql.Connection;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.Iterator;

import javax.sql.DataSource;

import org.postgresql.PGConnection;
import org.postgresql.PGNotification;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.netflix.conductor.core.exception.NonTransientException;
import com.netflix.conductor.postgres.config.PostgresProperties;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;

public class PostgresQueueListener {

private PGConnection pgconn;

private Connection conn;

private DataSource dataSource;

private HashMap<String, QueueStats> queues;

private boolean connected = false;

private boolean connecting = false;

private long lastNotificationTime = 0;

private Integer stalePeriod;

protected final Logger logger = LoggerFactory.getLogger(getClass());

public PostgresQueueListener(DataSource dataSource, PostgresProperties properties) {
logger.info("Using experimental PostgresQueueListener");
this.dataSource = dataSource;
this.stalePeriod = properties.getExperimentalQueueNotifyStalePeriod();
connect();
}

public boolean hasMessagesReady(String queueName) {
checkUpToDate();
handleNotifications();
if (notificationIsStale() || !connected) {
connect();
return true;
}

QueueStats queueStats = queues.get(queueName);
if (queueStats == null) {
return false;
}

if (queueStats.getNextDelivery() > System.currentTimeMillis()) {
return false;
}

return true;
}

public Integer getSize(String queueName) {
checkUpToDate();
handleNotifications();
if (notificationIsStale() || !connected) {
connect();
return null;
}

QueueStats queueStats = queues.get(queueName);
if (queueStats == null) {
return 0;
}

return queueStats.getDepth();
}

private boolean notificationIsStale() {
return System.currentTimeMillis() - lastNotificationTime > this.stalePeriod;
}

private void connect() {
if (!connecting) {
connected = false;
connecting = true;
try {
this.conn = dataSource.getConnection();
this.pgconn = conn.unwrap(PGConnection.class);

try {
boolean previousAutoCommitMode = conn.getAutoCommit();
conn.setAutoCommit(true);
try {
conn.prepareStatement("LISTEN conductor_queue_state").execute();
} catch (Throwable th) {
conn.rollback();
logger.info(th.getMessage());
} finally {
conn.setAutoCommit(previousAutoCommitMode);
connected = true;
}
} catch (SQLException ex) {
throw new NonTransientException(ex.getMessage(), ex);
}
requestStats();

} catch (SQLException e) {
logger.info(e.getSQLState());
}
connecting = false;
}
}

private void requestStats() {
try {
boolean previousAutoCommitMode = conn.getAutoCommit();
conn.setAutoCommit(true);
try {
conn.prepareStatement("SELECT queue_notify()").execute();
connected = true;
} catch (Throwable th) {
conn.rollback();
logger.info(th.getMessage());
} finally {
conn.setAutoCommit(previousAutoCommitMode);
}
} catch (SQLException e) {
if (e.getSQLState() != "08003") {
logger.info("Error fetching notifications {}", e.getSQLState());
}
connect();
}
}

private void checkUpToDate() {
if (System.currentTimeMillis() - lastNotificationTime > this.stalePeriod * 0.75) {
requestStats();
}
}

private void handleNotifications() {
try {
PGNotification[] notifications = pgconn.getNotifications();
if (notifications == null || notifications.length == 0) {
return;
}
processPayload(notifications[notifications.length - 1].getParameter());
} catch (SQLException e) {
if (e.getSQLState() != "08003") {
logger.info("Error fetching notifications {}", e.getSQLState());
}
connect();
}
}

private void processPayload(String payload) {
logger.info("Payload: {}", payload);
ObjectMapper objectMapper = new ObjectMapper();
try {
JsonNode notification = objectMapper.readTree(payload);
JsonNode lastNotificationTime = notification.get("__now__");
if (lastNotificationTime != null) {
this.lastNotificationTime = lastNotificationTime.asLong();
}
Iterator<String> iterator = notification.fieldNames();

HashMap<String, QueueStats> queueStats = new HashMap<>();
iterator.forEachRemaining(
key -> {
if (!key.equals("__now__")) {
QueueStats stats = null;
try {
stats =
objectMapper.treeToValue(
notification.get(key), QueueStats.class);
} catch (JsonProcessingException e) {
throw new RuntimeException(e);
}
queueStats.put(key, stats);
}
});
this.queues = queueStats;
logger.info(queues.toString());
} catch (JsonProcessingException e) {
throw new RuntimeException(e);
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
/*
* Copyright 2024 Conductor Authors.
* <p>
* Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
* <p>
* http://www.apache.org/licenses/LICENSE-2.0
* <p>
* Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
* an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
* specific language governing permissions and limitations under the License.
*/
package com.netflix.conductor.postgres.util;

public class QueueStats {
private Integer depth;

private long nextDelivery;

public void setDepth(Integer depth) {
this.depth = depth;
}

public Integer getDepth() {
return depth;
}

public void setNextDelivery(long nextDelivery) {
this.nextDelivery = nextDelivery;
}

public long getNextDelivery() {
return nextDelivery;
}

public String toString() {
return "{nextDelivery: " + nextDelivery + " depth: " + depth + "}";
}
}
Loading

0 comments on commit 0bf5a3d

Please sign in to comment.