external_ffmpeg/libavformat/cache.c

/*
 * Input cache protocol.
 * Copyright (c) 2011,2014 Michael Niedermayer
 *
 * This file is part of FFmpeg.
 *
 * FFmpeg is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * FFmpeg is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with FFmpeg; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
 *
 * Based on file.c by Fabrice Bellard
 */

/**
 * @TODO
 *      support keeping files
 *      support filling with a background thread
 */

#include "libavutil/avassert.h"
#include "libavutil/avstring.h"
#include "libavutil/internal.h"
#include "libavutil/opt.h"
#include "libavutil/tree.h"
#include "libavutil/time.h"
#include "avformat.h"
#include <pthread.h>
#include <fcntl.h>
#if HAVE_IO_H
#include <io.h>
#endif
#if HAVE_UNISTD_H
#include <unistd.h>
#endif
#include <sys/stat.h>
#include <stdlib.h>
#include <pthread.h>
#include "os_support.h"
#include "url.h"

#define CACHE_BUFSIZE   2048

typedef struct CacheEntry {
    int64_t logical_pos;
    int64_t physical_pos;
    int size;
} CacheEntry;

typedef struct Context {
    AVClass *class;
    int fd;
    char *filename;
    struct AVTreeNode *root;
    int64_t logical_pos;
    int64_t cache_pos;
    int64_t inner_pos;
    int64_t end;
    int is_true_eof;
    URLContext *inner;
    int64_t cache_hit, cache_miss;
    int read_ahead_limit;

    pthread_t thread;
    pthread_mutex_t mutex;

    int exit_request;
    AVIOInterruptCB interrupt_callback;
} Context;

static int cache_pread(URLContext *h, unsigned char *buf, int size, int64_t offset, bool complete);

static int cmp(const void *key, const void *node)
{
    return FFDIFFSIGN(*(const int64_t *)key, ((const CacheEntry *) node)->logical_pos);
}

static int cache_interrupt(void *arg)
{
    URLContext *h = arg;
    Context    *c = h->priv_data;

    if (c->exit_request)
        return 1;

    if (ff_check_interrupt(&c->interrupt_callback))
        c->exit_request = 1;

    return c->exit_request;
}

static void *cache_thread(void *arg)
{
    URLContext *h = arg;
    Context *c = h->priv_data;
    uint8_t buf[CACHE_BUFSIZE];
    int64_t time, bps, diff;
    int64_t offset = 0;
    int ret;

    time = av_gettime_relative();

    while (1) {
        pthread_mutex_lock(&c->mutex);
        if (cache_interrupt(h)) {
            ret = AVERROR_EXIT;
            pthread_mutex_unlock(&c->mutex);

            break;
        }
        if (offset < c->logical_pos)
            offset = c->logical_pos;

        ret = cache_pread(h, buf, sizeof(buf), offset, true);
        pthread_mutex_unlock(&c->mutex);

        if (ret > 0)
            offset += ret;
        else if (ret != AVERROR(EINTR))
            break;
    }

    diff = (av_gettime_relative() - time) / AV_TIME_BASE;
    bps = diff ? offset * 8 / diff : 0;
    av_log(h, AV_LOG_INFO, "%s end ret %d bytes %"PRId64" %"PRId64"bps\n", __func__, ret, offset, bps);

    return NULL;
}

static int cache_open(URLContext *h, const char *arg, int flags, AVDictionary **options)
{
    Context *c = h->priv_data;
    pthread_attr_t attr;
    ssize_t ssize;
    char *buffername;
    char threadname[128];
    int ret;
    AVIOInterruptCB interrupt_callback ={
        .callback = cache_interrupt,
        .opaque = h
    };

    av_strstart(arg, "cache:", &arg);

    c->fd = avpriv_tempfile("ffcache", &buffername, 0, h);
    if (c->fd < 0){
        av_log(h, AV_LOG_ERROR, "Failed to create tempfile %d\n", c->fd);
        return c->fd;
    }

    c->filename = buffername;

    c->interrupt_callback = h->interrupt_callback;
    ret = ffurl_open_whitelist(&c->inner, arg, flags, &interrupt_callback,
                                options, h->protocol_whitelist, h->protocol_blacklist, h);
    if (ret != 0) {
        av_log(h, AV_LOG_ERROR, "Failed to open: %s, %s\n", av_err2str(ret), arg);
        goto out;
    }

    pthread_mutex_init(&c->mutex, NULL);
    pthread_attr_init(&attr);

    ssize = pthread_get_stacksize_np(pthread_self());
    if (ssize > 0)
        pthread_attr_setstacksize(&attr, ssize);

    pthread_getname_np(pthread_self(), threadname, sizeof(threadname));
    av_strlcat(threadname, "_cache", sizeof(threadname));
    pthread_create(&c->thread, &attr, cache_thread, h);
    pthread_setname_np(c->thread, threadname);
    pthread_attr_destroy(&attr);

    return 0;

out:
    close(c->fd);
    if (c->filename) {
        unlink(c->filename);
        av_freep(&c->filename);
    }

    return ret;
}

static int add_entry(URLContext *h, const unsigned char *buf, int size, int64_t offset)
{
    Context *c= h->priv_data;
    int64_t pos = -1;
    int ret;
    CacheEntry *entry = NULL, *next[2] = {NULL, NULL};
    CacheEntry *entry_ret;
    struct AVTreeNode *node = NULL;

    //FIXME avoid lseek
    pos = lseek(c->fd, 0, SEEK_END);
    if (pos < 0) {
        ret = AVERROR(errno);
        av_log(h, AV_LOG_ERROR, "Failed to seek in cache, ret %"PRId64" errno %d\n", pos, errno);
        goto fail;
    }
    c->cache_pos = pos;

    ret = write(c->fd, buf, size);
    if (ret < 0) {
        av_log(h, AV_LOG_ERROR, "Failed to write %dB in cache, ret %d errno %d\n", size, ret, errno);
        ret = AVERROR(errno);
        goto fail;
    }
    c->cache_pos += ret;

    entry = av_tree_find(c->root, &offset, cmp, (void**)next);

    if (!entry)
        entry = next[0];

    if (!entry ||
        entry->logical_pos  + entry->size != offset ||
        entry->physical_pos + entry->size != pos
    ) {
        entry = av_malloc(sizeof(*entry));
        node = av_tree_node_alloc();
        if (!entry || !node) {
            ret = AVERROR(ENOMEM);
            goto fail;
        }
        entry->logical_pos = offset;
        entry->physical_pos = pos;
        entry->size = ret;

        entry_ret = av_tree_insert(&c->root, entry, cmp, &node);
        if (entry_ret && entry_ret != entry) {
            ret = -1;
            av_log(h, AV_LOG_ERROR, "av_tree_insert failed\n");
            goto fail;
        }
    } else
        entry->size += ret;

    return 0;
fail:
    //we could truncate the file to pos here if pos >=0 but ftruncate isn't available in VS so
    //for simplicty we just leave the file a bit larger
    av_free(entry);
    av_free(node);
    return ret;
}

static int cache_pread(URLContext *h, unsigned char *buf, int size, int64_t offset, bool complete)
{
    Context *c= h->priv_data;
    CacheEntry *entry, *next[2] = {NULL, NULL};
    int64_t r;

    entry = av_tree_find(c->root, &offset, cmp, (void**)next);
    if (!entry)
        entry = next[0];

    if (entry) {
        int64_t in_block_pos = offset - entry->logical_pos;
        av_assert0(entry->logical_pos <= offset);
        if (in_block_pos < entry->size) {
            int64_t physical_target = entry->physical_pos + in_block_pos;

            if (c->cache_pos != physical_target)
                r = lseek(c->fd, physical_target, SEEK_SET);
            else
                r = c->cache_pos;

            if (r >= 0) {
                c->cache_pos = r;
                r = read(c->fd, buf, FFMIN(size, entry->size - in_block_pos));
            }

            if (r > 0) {
                c->cache_pos += r;
                c->cache_hit ++;
                return r;
            }
        }
    }

    // Cache miss or some kind of fault with the cache

    if (offset != c->inner_pos) {
        r = ffurl_seek(c->inner, offset, SEEK_SET);
        if (r<0) {
            av_log(h, AV_LOG_ERROR, "Failed to perform internal seek\n");
            return r;
        }
        c->inner_pos = r;
    }

    if (complete)
        r = ffurl_read_complete(c->inner, buf, size);
    else
        r = ffurl_read(c->inner, buf, size);

    if (r == AVERROR_EOF && size>0) {
        c->is_true_eof = 1;
        av_assert0(c->end >= offset);
    }
    if (r<=0)
        return r;
    c->inner_pos += r;

    c->cache_miss ++;

    add_entry(h, buf, r, offset);
    c->end = FFMAX(c->end, offset + r);

    return r;
}

static int cache_read(URLContext *h, unsigned char *buf, int size)
{
    Context *c= h->priv_data;
    int r;

    r = cache_pread(h, buf, size, c->logical_pos, false);
    if (r > 0)
        c->logical_pos += r;

    return r;
}

static int cache_read_locked(URLContext *h, unsigned char *buf, int size)
{
    Context *c = h->priv_data;
    int ret;

    pthread_mutex_lock(&c->mutex);
    ret = cache_read(h, buf, size);
    pthread_mutex_unlock(&c->mutex);

    return ret;
}

static int64_t cache_seek(URLContext *h, int64_t pos, int whence)
{
    Context *c= h->priv_data;
    int64_t ret;

    if (whence == AVSEEK_SIZE) {
        pos= ffurl_seek(c->inner, pos, whence);
        if(pos <= 0){
            pos= ffurl_seek(c->inner, -1, SEEK_END);
            if (ffurl_seek(c->inner, c->inner_pos, SEEK_SET) < 0)
                av_log(h, AV_LOG_ERROR, "Inner protocol failed to seekback end : %"PRId64"\n", pos);
        }
        if (pos > 0)
            c->is_true_eof = 1;
        c->end = FFMAX(c->end, pos);
        return pos;
    }

    if (whence == SEEK_CUR) {
        whence = SEEK_SET;
        pos += c->logical_pos;
    } else if (whence == SEEK_END && c->is_true_eof) {
resolve_eof:
        whence = SEEK_SET;
        pos += c->end;
    }

    if (whence == SEEK_SET && pos >= 0 && pos < c->end) {
        //Seems within filesize, assume it will not fail.
        c->logical_pos = pos;
        return pos;
    }

    //cache miss
    ret= ffurl_seek(c->inner, pos, whence);
    if ((whence == SEEK_SET && pos >= c->logical_pos ||
         whence == SEEK_END && pos <= 0) && ret < 0) {
        if (   (whence == SEEK_SET && c->read_ahead_limit >= pos - c->logical_pos)
            || c->read_ahead_limit < 0) {
            uint8_t tmp[CACHE_BUFSIZE];
            while (c->logical_pos < pos || whence == SEEK_END) {
                int size = sizeof(tmp);
                if (whence == SEEK_SET)
                    size = FFMIN(sizeof(tmp), pos - c->logical_pos);
                ret = cache_read(h, tmp, size);
                if (ret == AVERROR_EOF && whence == SEEK_END) {
                    av_assert0(c->is_true_eof);
                    goto resolve_eof;
                }
                if (ret < 0) {
                    return ret;
                }
            }
            return c->logical_pos;
        }
    }

    if (ret >= 0) {
        c->logical_pos = ret;
        c->end = FFMAX(c->end, ret);
    }

    return ret;
}

static int64_t cache_seek_locked(URLContext *h, int64_t pos, int whence)
{
    Context *c = h->priv_data;
    int64_t ret;

    pthread_mutex_lock(&c->mutex);
    ret = cache_seek(h, pos, whence);
    pthread_mutex_unlock(&c->mutex);

    return ret;
}

static int enu_free(void *opaque, void *elem)
{
    av_free(elem);
    return 0;
}

static int cache_close(URLContext *h)
{
    Context *c= h->priv_data;
    int ret;

    av_log(h, AV_LOG_INFO, "Statistics, cache hits:%"PRId64" cache misses:%"PRId64"\n",
           c->cache_hit, c->cache_miss);

    pthread_mutex_lock(&c->mutex);
    c->exit_request = 1;
    pthread_mutex_unlock(&c->mutex);

    ret = pthread_join(c->thread, NULL);
    c->exit_request = 0;
    if (ret != 0)
        av_log(h, AV_LOG_ERROR, "Could not join thread %s.\n", av_err2str(ret));

    pthread_mutex_destroy(&c->mutex);
    close(c->fd);
    if (c->filename) {
        ret = unlink(c->filename);
        if (ret < 0)
            av_log(h, AV_LOG_ERROR, "Could not delete %s.\n", c->filename);
        av_freep(&c->filename);
    }
    ffurl_closep(&c->inner);
    av_tree_enumerate(c->root, NULL, NULL, enu_free);
    av_tree_destroy(c->root);

    return 0;
}

#define OFFSET(x) offsetof(Context, x)
#define D AV_OPT_FLAG_DECODING_PARAM

static const AVOption options[] = {
    { "read_ahead_limit", "Amount in bytes that may be read ahead when seeking isn't supported, -1 for unlimited", OFFSET(read_ahead_limit), AV_OPT_TYPE_INT, { .i64 = 65536 }, -1, INT_MAX, D },
    {NULL},
};

static const AVClass cache_context_class = {
    .class_name = "cache",
    .item_name  = av_default_item_name,
    .option     = options,
    .version    = LIBAVUTIL_VERSION_INT,
};

const URLProtocol ff_cache_protocol = {
    .name                = "cache",
    .url_open2           = cache_open,
    .url_read            = cache_read_locked,
    .url_seek            = cache_seek_locked,
    .url_close           = cache_close,
    .priv_data_size      = sizeof(Context),
    .priv_data_class     = &cache_context_class,
};