aos/usr/drivers/block_driver_server/main.c

#include <aos/aos.h>
#include <aos/paging.h>
#include <aos/ump_binding.h>
#include <aos/ump_chan.h>
#include <drivers/block_driver.h>
#include <aos/cache.h>

#include <maps/imx8x_map.h>
#include <drivers/sdhc.h>

// #define BLOCK_DRIVER_SERVER_PERFORMANCE
#ifdef BLOCK_DRIVER_SERVER_PERFORMANCE
#include <aos/performance.h>
static struct performance_context pcontext;
#endif

struct block_driver_state {
    struct ump_send_chan *send_chan;
    struct ump_recv_chan *recv_chan;
    struct block_driver_request *current_request;
    struct block_driver_state *next;
};

static struct sdhc_s *sdhc;

static void *dma_buffer;
static genpaddr_t dma_buffer_phys;

// a global lock for block driver access. requests waiting for the lock are put in the queue so they don't have to spin on the lock but locally on the ump channel instead
static bool locked = false;
static struct block_driver_state *lock_queue_head = NULL;
static struct block_driver_state *lock_queue_tail = NULL;

// list of open file handles because we have to synchronize on open dirs/files vs deletion
struct handle {
    uint64_t directory_entry_id;
    size_t count;
    struct handle *next;
    struct handle *prev;
};
struct handle *handle_list_head = NULL;
struct handle *handle_list_tail = NULL;

// very simple cache to keep the last block so it does not need to be reread all the time
// #define BLOCK_DRIVER_SINGLE_CACHE

#ifdef BLOCK_DRIVER_SINGLE_CACHE
uint32_t current_block_in_buffer = UINT32_MAX;
#endif

#define BLOCK_DRIVER_MULTI_CACHE
#ifdef BLOCK_DRIVER_MULTI_CACHE
#define BLOCK_DRIVER_CACHE_BLOCKS 2000
struct cache_metadata {
    uint32_t block_number;
    size_t access_time;
};
static char block_driver_cache[BLOCK_DRIVER_CACHE_BLOCKS * SDHC_BLOCK_SIZE];
static struct cache_metadata block_driver_cache_metadata[BLOCK_DRIVER_CACHE_BLOCKS];
static size_t block_driver_current_cache_time = 0;

static void init_cache(void) {
    for (size_t i = 0; i < BLOCK_DRIVER_CACHE_BLOCKS; ++i) {
        block_driver_cache_metadata[i].access_time = 0;
        block_driver_cache_metadata[i].block_number = UINT32_MAX;
    }
}

static void add_to_cache(uint32_t block_number) {
    assert(block_number != UINT32_MAX);
    // if it is already in the cache then update the access_time
    // also keep track of the entry with the lowest access_time
    size_t evict = 0;
    size_t lowest_time = SIZE_T_MAX;
    for (size_t i = 0; i < BLOCK_DRIVER_CACHE_BLOCKS; ++i) {
        if (block_driver_cache_metadata[i].block_number == block_number) {
            // debug_printf("Already in the cache %u at %lu\n", block_number, i);
            // update the value in the cache
            memcpy(&block_driver_cache[i * SDHC_BLOCK_SIZE], dma_buffer, SDHC_BLOCK_SIZE);
            block_driver_cache_metadata[i].access_time = ++block_driver_current_cache_time;
            return;
        }
        if (block_driver_cache_metadata[i].access_time < lowest_time) {
            lowest_time = block_driver_cache_metadata[i].access_time;
            evict = i;
        }
    }
    
    // debug_printf("Adding %u to cache at %lu\n", block_number, evict);

    // add it to the cache and evict an entry with the smallest access_time (oldest/LRU)
    memcpy(&block_driver_cache[evict * SDHC_BLOCK_SIZE], dma_buffer, SDHC_BLOCK_SIZE);
    block_driver_cache_metadata[evict].access_time = ++block_driver_current_cache_time;
    block_driver_cache_metadata[evict].block_number = block_number;
}

static bool load_from_cache(uint32_t block_number) {
    assert(block_number != UINT32_MAX);
    for (size_t i = 0; i < BLOCK_DRIVER_CACHE_BLOCKS; ++i) {
        if (block_driver_cache_metadata[i].block_number == block_number) {
            // debug_printf("HIT for %u at %lu\n", block_number, i);
            block_driver_cache_metadata[i].access_time = ++block_driver_current_cache_time;
            memcpy(dma_buffer, &block_driver_cache[i * SDHC_BLOCK_SIZE], SDHC_BLOCK_SIZE);
            __asm volatile (
                "dmb sy\n"
            );
            cpu_dcache_wbinv_range((genvaddr_t)dma_buffer, SDHC_BLOCK_SIZE);
            return true;
        }
    }
    // debug_printf("MISS!\n");
    return false;
}

#endif

// don't read a block during writing if the whole block gets overwritten
#define BLOCK_DRIVER_OPTIMIZE_FULL_WRITE

static errval_t read_buffer(int block_number) {
    assert(block_number != UINT32_MAX);

    errval_t err;

    // debug_printf("[block_driver_server] read block %d\n", block_number);

    #ifdef BLOCK_DRIVER_SINGLE_CACHE
    // very primitive cache for the last read block
    if(current_block_in_buffer == block_number) {
        return SYS_ERR_OK;
    }
    #endif

    #ifdef BLOCK_DRIVER_MULTI_CACHE
    if(load_from_cache(block_number)) {
        return SYS_ERR_OK;   
    }
    #endif

    err = sdhc_read_block(sdhc, block_number, (lpaddr_t)dma_buffer_phys);
    if(err_is_fail(err)) {
        #ifdef BLOCK_DRIVER_SINGLE_CACHE
        // maybe there was a partial modification so invalidate the cache
        current_block_in_buffer = UINT32_MAX;
        #endif
        // debug_printf("[block_driver_server] err done\n");
        return err;
    }
    __asm volatile (
        "dmb sy\n"
    );
    cpu_dcache_wbinv_range((genvaddr_t)dma_buffer, SDHC_BLOCK_SIZE);

    #ifdef BLOCK_DRIVER_SINGLE_CACHE
    current_block_in_buffer = block_number;
    #endif
    #ifdef BLOCK_DRIVER_MULTI_CACHE
    add_to_cache(block_number);
    #endif

    return SYS_ERR_OK;
}

static errval_t write_buffer(int block_number) {
    assert(block_number != UINT32_MAX);

    errval_t err;

    debug_printf("[block_driver_server] write block %d\n", block_number);

    __asm volatile (
        "dmb sy\n"
    );
    cpu_dcache_wbinv_range((genvaddr_t)dma_buffer, SDHC_BLOCK_SIZE);

    err = sdhc_write_block(sdhc, block_number, (lpaddr_t)dma_buffer_phys);
    if(err_is_fail(err)) return err;

    // debug_printf("[block_driver_server] done\n");
    return SYS_ERR_OK;
}

/**
 * @brief search for open handles for the given directory entry id
 */
static struct handle *search_element(uint64_t directory_entry_id) {
    struct handle *cur_handle = handle_list_head;
    while(cur_handle != NULL && cur_handle->directory_entry_id != directory_entry_id) {
        cur_handle = cur_handle->next;
    }
    return cur_handle;
}

/**
 * @brief insert a structure in the list of directory entry handles
 */
static void insert_element(uint64_t directory_entry_id) {
    struct handle *element_handle = calloc(1, sizeof(struct handle));
    if (element_handle == NULL) USER_PANIC("Ran out of memory");

    element_handle->count = 1;
    element_handle->directory_entry_id = directory_entry_id;

    element_handle->prev = handle_list_tail;
    element_handle->next = NULL;
    if(handle_list_head == NULL) {
        handle_list_head = element_handle;
    } else {
        handle_list_tail->next = element_handle;
    }
    handle_list_tail = element_handle;
}

/**
 * @brief remove an element from the list of directory handles
 */
static void remove_element(struct handle *element_handle) {
    if (element_handle->prev == NULL) {
        handle_list_head = element_handle->next;
    } else {
        element_handle->prev->next = element_handle->next;
    }
    if (element_handle->next == NULL) {
        handle_list_tail = element_handle->prev;
    } else {
        element_handle->next->prev = element_handle->prev;
    }
}

__attribute__((__unused__))
/**
 * @brief print all open handles for debugging
 */
static void print_handles(void) {
    struct handle *element_handle = handle_list_head;
    debug_printf("Handle List\n");
    while (element_handle != NULL) {
        debug_printf("Handle %lu: %lu\n", element_handle->directory_entry_id, element_handle->count);
        element_handle = element_handle->next;
    }
}

/**
 * @brief add an open handle for the given directory entry id
 */
static errval_t register_handle(uint64_t directory_entry_id) {
    // debug_printf("Register %lu\n", directory_entry_id);
    struct handle *element_handle = search_element(directory_entry_id);
    if (element_handle == NULL) {
        insert_element(directory_entry_id);
    } else if(element_handle->count == SIZE_MAX) {
        return FAT_ERR_TOO_MANY_HANDLES;
    }else {
        ++element_handle->count;
    }
    // print_handles();

    return SYS_ERR_OK;
}

/**
 * @brief remove an open handle from the given directory entry id
 */
static errval_t unregister_handle(uint64_t directory_entry_id) {
    // debug_printf("Unegister %lu\n", directory_entry_id);
    struct handle *element_handle = search_element(directory_entry_id);
    if (element_handle == NULL) {
        return ERR_INVALID_ARGS;
    } else if (element_handle->count == 1) {
        remove_element(element_handle);
    } else {
        --element_handle->count;
    }
    // print_handles();
    return SYS_ERR_OK;
}

static size_t count_handles(uint64_t directory_entry_id) {
    struct handle *element_handle = search_element(directory_entry_id);
    if (element_handle == NULL) {
        return 0;
    } else {
        return element_handle->count;
    }
}

/**
 * @brief read an object from the sdcard
 */
static errval_t read_object(uint32_t block_number, size_t offset, size_t size, void *dst) {
    errval_t err;

    // printf("[block_driver_server] reading %lu bytes from %lu@%u to %p\n", size, offset, block_number, dst);

    // sector bounds check
    if (offset + size > SDHC_BLOCK_SIZE) {
        return BLOCK_ERR_OVERFLOW_BLOCK;
    }

    #ifdef BLOCK_DRIVER_SERVER_PERFORMANCE
    perf_init(&pcontext, "block_driver_read_object");
    perf_add_now(&pcontext, "start");
    #endif

    err = read_buffer(block_number);
    if (err_is_fail(err)) return err;

    #ifdef BLOCK_DRIVER_SERVER_PERFORMANCE
    perf_add_now(&pcontext, "memcpy");
    #endif

    memcpy(dst, dma_buffer + offset, size);

    #ifdef BLOCK_DRIVER_SERVER_PERFORMANCE
    perf_add_now(&pcontext, "done");
    perf_print(&pcontext);
    #endif

    return SYS_ERR_OK;
}

/**
 * @brief write an object to the sdcard
 */
static errval_t write_object(uint32_t block_number, size_t offset, size_t size, void *src) {
    errval_t err;

    // printf("[block_driver_server] writing %lu bytes from %p to %lu@%u\n", size, src, offset, block_number);

    // sector bounds check
    if (offset + size > SDHC_BLOCK_SIZE) {
        return BLOCK_ERR_OVERFLOW_BLOCK;
    }

    // we never modify the bpb and boot sector
    if (block_number == 0) {
        return BLOCK_ERR_WRITE_BOOTSECTOR;
    }

    #ifdef BLOCK_DRIVER_SERVER_PERFORMANCE
    perf_init(&pcontext, "block_driver_write_object");
    perf_add_now(&pcontext, "start");
    #endif

    // if we are not writing the full block then we need to copy the current block
    #ifdef BLOCK_DRIVER_OPTIMIZE_FULL_WRITE
    if (size < SDHC_BLOCK_SIZE) {
    #endif
        err = read_buffer(block_number);
        if (err_is_fail(err)) return err;
    #ifdef BLOCK_DRIVER_OPTIMIZE_FULL_WRITE
    }
    #endif

    #ifdef BLOCK_DRIVER_SERVER_PERFORMANCE
    perf_add_now(&pcontext, "read");
    #endif

    // debug_printf("Write start: '%.5s'\n", src);
    memcpy(dma_buffer + offset, src, size);
    
    #ifdef BLOCK_DRIVER_SINGLE_CACHE
    // update the cache metadata
    current_block_in_buffer = block_number;
    #endif
    #ifdef BLOCK_DRIVER_MULTI_CACHE
    add_to_cache(block_number);
    #endif

    #ifdef BLOCK_DRIVER_SERVER_PERFORMANCE
    perf_add_now(&pcontext, "memcpy");
    #endif

    err = write_buffer(block_number);
    if (err_is_fail(err)) return err;

    #ifdef BLOCK_DRIVER_SERVER_PERFORMANCE
    perf_add_now(&pcontext, "done");
    perf_print(&pcontext);
    #endif

    return SYS_ERR_OK;
}

static void handle_send_completed(void *arg, struct ump_send_queue_entry *entry) {
    free((void *)entry->header);
    free((void *)entry->payload);
    free(entry);
}

static void send_response(struct block_driver_state *state, errval_t err, size_t payload_size, void *payload) {
    assert(state != NULL);

    struct block_driver_result *result = malloc(sizeof(struct block_driver_result));
    if (result == NULL) USER_PANIC("Ran out of memory");
    struct ump_send_queue_entry *entry = malloc(sizeof(struct ump_send_queue_entry));
    if (entry == NULL) USER_PANIC("Ran out of memory");

    result->err = err;
    ump_send(
        state->send_chan,
        entry,
        sizeof(struct block_driver_result),
        result,
        payload_size,
        payload,
        handle_send_completed,
        NULL
    );
}
static void send_error_response(struct block_driver_state *state, errval_t err) {
    assert(state != NULL);

    send_response(state, err, 0, NULL);
}

static void handle_payload(void *arg, size_t payload_size, void *payload);
static void handle_request(void *arg, size_t header_size, void *header, size_t payload_size) {
    assert(arg != NULL);

    errval_t err;

    struct block_driver_state *state = arg;
    struct block_driver_request *request = header;

    // keep track of which request we are currently handling for a given connection
    state->current_request = request;

    if (request->bytes > SDHC_BLOCK_SIZE) {
        send_error_response(state, ERR_INVALID_ARGS);
    } else if (request->action == BLOCK_DRIVER_ACTION_READ) {
        // if no bytes were requested just return ok. This allows measuring performance of this UMP interface
        if(request->bytes == 0) {
            send_response(state, SYS_ERR_OK, 0, NULL);
        } else {
            // perform read directly
            void *response_buffer = malloc(request->bytes);
            if (response_buffer == NULL) USER_PANIC("Ran out of memory");

            err = read_object(request->block_number, request->offset, request->bytes, response_buffer);
            if (err_is_fail(err)) {
                send_error_response(state, err);
            } else {
                send_response(state, SYS_ERR_OK, request->bytes, response_buffer);
            }
        }
    } else if (request->action == BLOCK_DRIVER_ACTION_WRITE) {
        if (request->bytes != payload_size) {
            send_error_response(state, ERR_INVALID_ARGS);
            ump_recv_payload(state->recv_chan, NULL, handle_payload, arg);
            return;
        }
        // we first need to receive the payload before we can handle the request
        void *payload_buffer = malloc(request->bytes);
        if (payload_buffer == NULL) USER_PANIC("Ran out of memory");
        ump_recv_payload(state->recv_chan, payload_buffer, handle_payload, arg);
        return;
    } else if (request->action == BLOCK_DRIVER_ACTION_LOCK) {
        if(locked) {
            // enqueue the request
            if (lock_queue_tail != NULL) {
                lock_queue_tail->next = state;
            } else {
                lock_queue_head = state;
            }
            lock_queue_tail = state;
        } else {
            locked = true;
            send_error_response(state, SYS_ERR_OK);
        }
    } else if (request->action == BLOCK_DRIVER_ACTION_UNLOCK) {
        if (!locked) USER_PANIC("Received unlock command but was not locked");
        send_error_response(state, SYS_ERR_OK);

        // send response to first thread on the lock queue if there is one
        if (lock_queue_head != NULL) {
            send_error_response(lock_queue_head, SYS_ERR_OK);

            // dequeue the head
            if (lock_queue_head->next == NULL) {
                lock_queue_tail = NULL;
            }
            lock_queue_head = lock_queue_head->next;
        } else {
            locked = false;
        }
    } else if (request->action == BLOCK_DRIVER_ACTION_REGISTER_HANDLE) {
        err = register_handle(request->directory_entry_id);
        send_error_response(state, err);
    } else if (request->action == BLOCK_DRIVER_ACTION_UNREGISTER_HANDLE) {
        err = unregister_handle(request->directory_entry_id);
        send_error_response(state, err);
    } else if (request->action == BLOCK_DRIVER_ACTION_COUNT_HANDLES) {
        size_t *count = malloc(sizeof(size_t));
        if (count == NULL) {
            send_error_response(state, LIB_ERR_MALLOC_FAIL);
        } else {
            *count = count_handles(request->directory_entry_id);
            send_response(state, SYS_ERR_OK, sizeof(size_t), count);
        }
    } else {
        send_error_response(state, ERR_INVALID_ARGS);
    }

    // skip payload
    ump_recv_payload(state->recv_chan, NULL, handle_payload, arg);
}

static void handle_payload(void *arg, size_t payload_size, void *payload) {
    assert(arg != NULL);

    errval_t err;

    struct block_driver_state *state = arg;
    struct block_driver_request *request = state->current_request;

    // we are guaranteed to have a valid structure request here
    if (request->action == BLOCK_DRIVER_ACTION_WRITE) {
        // we got the payload, write it to disk
        err = write_object(request->block_number, request->offset, request->bytes, payload);
        send_response(state, err, 0, NULL);
        // free the payload buffer, we don't use it anymore
        free(payload);
    }

    // listen for the next request on this channel
    ump_recv_header(state->recv_chan, handle_request, arg);
}

static errval_t connect_callback(void *arg, struct capref cap) {
    errval_t err;

    // printf("[block_driver_server] incoming connection\n");

    struct ump_send_chan *send_chan;
    struct ump_recv_chan *recv_chan;

    // we can run the server on the default waitset since we are dispatching on
    // it for listening to connections anyways
    err = ump_chan_init(UMP_ROLE_SERVER, &send_chan, &recv_chan, sizeof(struct block_driver_request), cap, get_default_waitset());
    if (err_is_fail(err)) return err;

    struct block_driver_state *state = malloc(sizeof(struct block_driver_state));
    if (state == NULL) return LIB_ERR_MALLOC_FAIL;
    state->recv_chan = recv_chan;
    state->send_chan = send_chan;
    ump_recv_header(recv_chan, handle_request, state);

    // printf("[block_driver_server] connection ready to receive requests\n");

    return SYS_ERR_OK;
}

int main(int argc, char *argv[]) {

    errval_t err;

    // get the device capability from the arg cnode
    struct capref cap_arg0 = {
        .cnode = cnode_arg,
        .slot  = 0
    };
    // map device registers uncachable
    void *device_register_vaddr;
    err = paging_map_frame_attr(
        get_current_paging_state(),
        &device_register_vaddr,
        IMX8X_SDHC_SIZE,
        cap_arg0,
        VREGION_FLAGS_READ_WRITE_NOCACHE
    );
    if (err_is_fail(err)) USER_PANIC_ERR(err, "Failed to map device registers");

    // initialize the sdhc driver
    DEBUG_PRINTF("[block_driver_server] initializing sdhc driver\n");
    err = sdhc_init(&sdhc, device_register_vaddr);
    if (err_is_fail(err)) USER_PANIC_ERR(err, "Failed to initialize sdhc driver");

    // create the dma buffer
    struct capref dma_buffer_frame;
    size_t size;
    err = frame_alloc(&dma_buffer_frame, SDHC_BLOCK_SIZE, &size);
    if(err_is_fail(err)) return err_push(err, LIB_ERR_FRAME_ALLOC);

    // we need the physical address for the DMA read operation
    struct frame_identity dma_buffer_frame_id;
    err = frame_identify(dma_buffer_frame, &dma_buffer_frame_id);
    if (err_is_fail(err)) return err_push(err, LIB_ERR_FRAME_IDENTIFY);
    dma_buffer_phys = dma_buffer_frame_id.base;

    // map the scratchpad frame for us
    err = paging_map_frame_attr(get_current_paging_state(), &dma_buffer, size, dma_buffer_frame, VREGION_FLAGS_READ_WRITE);
    if (err_is_fail(err)) return err_push(err, LIB_ERR_PMAP_MAP);

    struct ump_binding_server server;
    err = ump_binding_register(&server, UMP_SERVER_BLOCK_DRIVER, connect_callback, NULL);
    if (err_is_fail(err)) USER_PANIC_ERR(err, "Failed to register UMP server");

    #ifdef BLOCK_DRIVER_MULTI_CACHE
    // initialize the cache
    init_cache();
    #endif

    // wait for connections to the server
    printf("[block_driver_server] ready\n");
    struct waitset *default_ws = get_default_waitset();
    while (true) {
        err = event_dispatch(default_ws);
        if (err_is_fail(err)) {
            DEBUG_ERR(err, "in event_dispatch");
            abort();
        }
    }

    return EXIT_SUCCESS;
}