task_scheduler.c

/*
 * Copyright 2010-2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License").
 * You may not use this file except in compliance with the License.
 * A copy of the License is located at
 *
 *  http://aws.amazon.com/apache2.0
 *
 * or in the "license" file accompanying this file. This file 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.
 */

#include <aws/common/task_scheduler.h>
#include <aws/common/thread.h>
#include <aws/common/priority_queue.h>

static const size_t DEFAULT_QUEUE_SIZE = 7;

struct task_container {
    uint64_t timestamp;
    struct aws_task task;
};

static int compare_timestamps(const void *a, const void *b) {
    uint64_t a_time = ((struct task_container*)a)->timestamp;
    uint64_t b_time = ((struct task_container*)b)->timestamp;
    return a_time > b_time; /* min-heap */
}

static inline int get_next_task(struct aws_task_scheduler *scheduler, struct aws_task *task, uint64_t run_before,
                                uint64_t *next_run_time);

int aws_task_scheduler_init(struct aws_task_scheduler *scheduler, struct aws_allocator *alloc, 
        aws_task_scheduler_clock clock) {
    scheduler->alloc = alloc;
    scheduler->clock = clock;
    scheduler->min_run_time = 0;
    return aws_priority_queue_dynamic_init(&scheduler->queue, alloc, DEFAULT_QUEUE_SIZE, sizeof(struct task_container),
            &compare_timestamps);
}

void aws_task_scheduler_clean_up(struct aws_task_scheduler *scheduler) {
    uint64_t everything_in_past = UINT64_MAX;

    while(1) {
        struct aws_task task_to_run = {0};

        if (get_next_task(scheduler, &task_to_run, everything_in_past, NULL)) {
            break;
        }

        task_to_run.fn(task_to_run.arg, AWS_TASK_STATUS_CANCELED);
    }

    aws_priority_queue_clean_up(&scheduler->queue);
    scheduler->alloc = NULL;
    scheduler->clock = NULL;
}

static inline int get_next_task(struct aws_task_scheduler *scheduler, struct aws_task *task, uint64_t run_before,
                                uint64_t *next_run_time) {
    struct task_container *possible_task;
    if(aws_priority_queue_top(&scheduler->queue, (void **)&possible_task)) {
        if(AWS_ERROR_PRIORITY_QUEUE_EMPTY == aws_last_error()) {
            return aws_raise_error(AWS_ERROR_TASK_SCHEDULER_NO_TASKS);
        }
        return AWS_OP_ERR;
    }

    if (next_run_time) {
        *next_run_time = possible_task->timestamp;
    }

    if (possible_task->timestamp > run_before) {
        return aws_raise_error(AWS_ERROR_TASK_SCHEDULER_NO_READY_TASKS);
    }

    struct task_container container;
    if(aws_priority_queue_pop(&scheduler->queue, (void *)&container)) {
        return AWS_OP_ERR;
    }

    *task = container.task;
    if (next_run_time) {
        if(aws_priority_queue_top(&scheduler->queue, (void **)&possible_task)) {
            *next_run_time = 0;
            if(AWS_ERROR_PRIORITY_QUEUE_EMPTY == aws_last_error()) {
                return AWS_OP_SUCCESS;
            }
            return AWS_OP_ERR;
        }
        *next_run_time = possible_task->timestamp;
    }
    return AWS_OP_SUCCESS;
}

int aws_task_scheduler_next_task(struct aws_task_scheduler *scheduler, 
        struct aws_task *task, uint64_t *next_run_time) {

    uint64_t now;
    if(scheduler->clock(&now)) {
        return AWS_OP_ERR;
    }

    return get_next_task(scheduler, task, now, next_run_time);
}

int aws_task_scheduler_schedule_now(struct aws_task_scheduler *scheduler,
        struct aws_task *task) {

    uint64_t now;
    if(scheduler->clock(&now)) {
        return AWS_OP_ERR;
    }

    return aws_task_scheduler_schedule_future(scheduler, task, now);
}

int aws_task_scheduler_schedule_future(struct aws_task_scheduler *scheduler, 
        struct aws_task *task, uint64_t time_to_run) {

    struct task_container container;

    /* this serves the purpose of handling reentrant scheduling while the tasks are being run.
       if the clock tick is on the same nanosecond (with microsecond precision) as the run was kicked
       off, it will still be run. As a result, if time_to_run is before or on the same clock tick, increment it
       by one nanosecond to avoid the issue. */
    if (AWS_UNLIKELY(time_to_run < scheduler->min_run_time)) {
        time_to_run = scheduler->min_run_time;
    }

    container.task = *task;
    container.timestamp = time_to_run;

    if(aws_priority_queue_push(&scheduler->queue, &container)) {
        return AWS_OP_ERR;
    }

    return AWS_OP_SUCCESS;
}

int aws_task_scheduler_run_all(struct aws_task_scheduler *scheduler, uint64_t *next_task_time) {

    uint64_t now;
    if(scheduler->clock(&now)) {
        return AWS_OP_ERR;
    }

    scheduler->min_run_time = now + 1;
    while(1) {
        struct aws_task task_to_run = {0};

        if (get_next_task(scheduler, &task_to_run,  now, next_task_time)) {
            int err_code = aws_last_error();
            if (err_code == AWS_ERROR_TASK_SCHEDULER_NO_READY_TASKS  || err_code == AWS_ERROR_TASK_SCHEDULER_NO_TASKS) {
                return AWS_OP_SUCCESS;
            }

            return AWS_OP_ERR;
        }

        task_to_run.fn(task_to_run.arg, AWS_TASK_STATUS_RUN_READY);
    }
}