io.pythagoras.common.aggregationqueue.QueueManager Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of aggregation-queue Show documentation
The newest version!
package io.pythagoras.common.aggregationqueue;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import javax.persistence.EntityManager;
import javax.persistence.NoResultException;
import javax.persistence.Query;
import java.util.Date;
import java.util.UUID;

@Service
public class QueueManager {

    private EntityManager entityManager;

    private int offset = 1800; // Default 30 min

    @Autowired
    public void setEntityManager(EntityManager entityManager) {
        this.entityManager = entityManager;
    }

    public void setOffset(int offset) {
        this.offset = offset;
    }

    public Task insertTask(String code, String taskType) {

        Query query = entityManager.createNativeQuery("INSERT INTO ipcag_task (code,task_type, timestmp, randomizer, request_count, execution_count) " +
                "VALUES (:code, :taskType, :timestmp, :randomizer, 1, 0) ON CONFLICT ON CONSTRAINT code_timestmp " +
                "DO UPDATE SET randomizer = :randomizer , request_count = ( ipcag_task.request_count + 1 ) RETURNING *", Task.class);

        query.setParameter("code", code);
        query.setParameter("taskType", taskType);
        query.setParameter("timestmp", new Date(0));
        query.setParameter("randomizer", UUID.randomUUID().toString());

        Task result = (Task) query.getSingleResult();
        return result;
    }

    public Task claimNextTask() {
        return processClaimNextTask(makeQueryForClaimAnyTask());
    }

    public Task claimNextTaskByType(String taskType) {
        return processClaimNextTask(makeQueryForClaimTaskByType(taskType));
    }

    Task processClaimNextTask(Query query) {
        try {
            Task result = (Task) query.getSingleResult();
            return result;
        } catch (NoResultException e) {
            return null;
        }
    }

    Query makeQueryForClaimAnyTask() {
        return entityManager.createNativeQuery("UPDATE ipcag_task SET timestmp = now(), " +
                "execution_count = (execution_count + 1) " +
                "WHERE id = ( " +
                "SELECT id FROM ipcag_task " +
                "WHERE timestmp < CAST ((now() - INTERVAL '" + offset + " second' ) as TIMESTAMP ) " +
                "AND code NOT IN ( " +
                "SELECT code FROM ipcag_task " +
                "WHERE timestmp >= CAST ((now() - INTERVAL '" + offset + " second' ) as TIMESTAMP )) LIMIT 1 ) " +
                "RETURNING *", Task.class);
    }

    Query makeQueryForClaimTaskByType(String taskType) {
        Query query = entityManager.createNativeQuery("UPDATE ipcag_task SET timestmp = now(), " +
                "execution_count = (execution_count + 1) " +
                "WHERE id = ( " +
                "SELECT id FROM ipcag_task " +
                "WHERE timestmp < CAST ((now() - INTERVAL '" + offset + " second' ) as TIMESTAMP ) " +
                "AND task_type = :taskType " +
                "AND code NOT IN ( " +
                "SELECT code FROM ipcag_task " +
                "WHERE timestmp >= CAST ((now() - INTERVAL '" + offset + " second' ) as TIMESTAMP )) LIMIT 1 ) " +
                "RETURNING *", Task.class);

        query.setParameter("taskType", taskType);
        return query;
    }

    public void completeTask(Task task) throws AggregationQueueException {
        completeTask(task.getId());

    }

    public void completeTask(int taskId) throws AggregationQueueException {
        Query query = entityManager.createNativeQuery("DELETE FROM ipcag_task WHERE id = :id RETURNING id");

        query.setParameter("id", taskId);

        try {
            Integer result = (Integer) query.getSingleResult();
            if(!result.equals(taskId)) {
                throw new AggregationQueueException("Completed unexpected job.  Expected: "+taskId+" Completed: "+result);
            }
        } catch (NoResultException e) {
            throw new AggregationQueueException("Unable to complete job: " + taskId,e);
        }


    }

}