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 23, 2024
1 parent 06bc5fd commit c9d7bf4
Show file tree
Hide file tree
Showing 8 changed files with 522 additions and 4 deletions.
2 changes: 1 addition & 1 deletion dependencies.gradle
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ ext {
revActivation = '2.0.1'
revApacheHttpComponentsClient5 = '5.2.1'
revAwaitility = '3.1.6'
revAwsSdk = '1.12.635'
revAwsSdk = '1.11.86'
revBval = '2.0.5'
revCassandra = '3.10.2'
revCassandraUnit = '3.11.2.0'
Expand Down
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,8 +27,20 @@ public class PostgresProperties {

private Integer deadlockRetryMax = 3;

private boolean experimentalQueueNotify = false;

private Integer experimentalQueueNotifyStalePeriod = 5000;

public String schema = "public";

public Integer getExperimentalQueueNotifyStalePeriod() {
return experimentalQueueNotifyStalePeriod;
}

public void setExperimentalQueueNotifyStalePeriod(Integer experimentalQueueNotifyStalePeriod) {
this.experimentalQueueNotifyStalePeriod = experimentalQueueNotifyStalePeriod;
}

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

public boolean getExperimentalQueueNotify() {
return experimentalQueueNotify;
}

public void setExperimentalQueueNotify(boolean experimentalQueueNotify) {
this.experimentalQueueNotify = experimentalQueueNotify;
}
}
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,189 @@
/*
* 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 com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.netflix.conductor.postgres.config.PostgresProperties;
import org.postgresql.PGConnection;
import org.postgresql.PGNotification;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.netflix.conductor.core.exception.NonTransientException;

public class PostgresQueueListener {

private PGConnection pgconn;

private Connection conn;

private DataSource dataSource;

private HashMap<String, QueueStats> queues;

private boolean connected = false;

private long lastNotificationTime = 0;

private Integer stalePeriod;

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

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

public boolean hasMessagesReady(String queueName) {
checkUpToDate();
handleNotifications();
if (System.currentTimeMillis() - lastNotificationTime > 6000){
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 (System.currentTimeMillis() - lastNotificationTime > 6000){
return null;
}

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

return queueStats.getDepth();
}

private void connect(){
connected = false;
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; SELECT queue_notify();").execute();
connected = true;
} catch (Throwable th) {
conn.rollback();
logger.info(th.getMessage());
} finally {
conn.setAutoCommit(previousAutoCommitMode);
}
} catch (SQLException ex) {
throw new NonTransientException(ex.getMessage(), ex);
}
requestStats();

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

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 ex) {
throw new NonTransientException(ex.getMessage(), ex);
}
}

private void checkUpToDate(){
if (System.currentTimeMillis() - lastNotificationTime > this.stalePeriod){
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) {
logger.info("Error fetching notifications {}", e.getSQLState());
if(e.getSQLState() == "08003"){
logger.info("Reconnecting");
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,27 @@
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 + "}";
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
CREATE OR REPLACE FUNCTION queue_notify() RETURNS void
LANGUAGE SQL
AS $$
SELECT pg_notify('conductor_queue_state', (
SELECT
COALESCE(jsonb_object_agg(KEY, val), '{}'::jsonb) ||
jsonb_build_object('__now__', (extract('epoch' from CURRENT_TIMESTAMP)*1000)::bigint)
FROM (
SELECT
queue_name AS KEY,
jsonb_build_object(
'nextDelivery',
(extract('epoch' from min(deliver_on))*1000)::bigint,
'depth',
count(*)
) AS val
FROM
queue_message
WHERE
popped = FALSE
GROUP BY
queue_name) AS sq)::text);
$$;


CREATE FUNCTION queue_notify_trigger()
RETURNS TRIGGER
LANGUAGE PLPGSQL
AS $$
BEGIN
PERFORM queue_notify();
RETURN NULL;
END;
$$;

CREATE TRIGGER queue_update
AFTER UPDATE ON queue_message
FOR EACH ROW
WHEN (OLD.popped IS DISTINCT FROM NEW.popped)
EXECUTE FUNCTION queue_notify_trigger();

CREATE TRIGGER queue_insert_delete
AFTER INSERT OR DELETE ON queue_message
FOR EACH ROW
EXECUTE FUNCTION queue_notify_trigger();

CREATE TRIGGER queue_delete
AFTER DELETE ON queue_message
FOR EACH ROW
EXECUTE FUNCTION queue_notify_trigger();
Loading

0 comments on commit c9d7bf4

Please sign in to comment.