fftools/thread_queue.c

   1 /*
   2  * This file is part of FFmpeg.
   3  *
   4  * FFmpeg is free software; you can redistribute it and/or
   5  * modify it under the terms of the GNU Lesser General Public
   6  * License as published by the Free Software Foundation; either
   7  * version 2.1 of the License, or (at your option) any later version.
   8  *
   9  * FFmpeg is distributed in the hope that it will be useful,
  10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  12  * Lesser General Public License for more details.
  13  *
  14  * You should have received a copy of the GNU Lesser General Public
  15  * License along with FFmpeg; if not, write to the Free Software
  16  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  17  */
  18
  19 #include <stdint.h>
  20 #include <string.h>
  21
  22 #include "libavutil/avassert.h"
  23 #include "libavutil/error.h"
  24 #include "libavutil/fifo.h"
  25 #include "libavutil/intreadwrite.h"
  26 #include "libavutil/mem.h"
  27 #include "libavutil/thread.h"
  28
  29 #include "objpool.h"
  30 #include "thread_queue.h"
  31
  32 enum {
  33     FINISHED_SEND = (1 << 0),
  34     FINISHED_RECV = (1 << 1),
  35 };
  36
  37 typedef struct FifoElem {
  38     void        *obj;
  39     unsigned int stream_idx;
  40 } FifoElem;
  41
  42 struct ThreadQueue {
  43     int              *finished;
  44     unsigned int    nb_streams;
  45
  46     AVFifo  *fifo;
  47
  48     ObjPool *obj_pool;
  49     void   (*obj_move)(void *dst, void *src);
  50
  51     pthread_mutex_t lock;
  52     pthread_cond_t  cond;
  53 };
  54
  55 void tq_free(ThreadQueue **ptq)
  56 {
  57     ThreadQueue *tq = *ptq;
  58
  59     if (!tq)
  60         return;
  61
  62     if (tq->fifo) {
  63         FifoElem elem;
  64         while (av_fifo_read(tq->fifo, &elem, 1) >= 0)
  65             objpool_release(tq->obj_pool, &elem.obj);
  66     }
  67     av_fifo_freep2(&tq->fifo);
  68
  69     objpool_free(&tq->obj_pool);
  70
  71     av_freep(&tq->finished);
  72
  73     pthread_cond_destroy(&tq->cond);
  74     pthread_mutex_destroy(&tq->lock);
  75
  76     av_freep(ptq);
  77 }
  78
  79 ThreadQueue *tq_alloc(unsigned int nb_streams, size_t queue_size,
  80                       ObjPool *obj_pool, void (*obj_move)(void *dst, void *src))
  81 {
  82     ThreadQueue *tq;
  83     int ret;
  84
  85     tq = av_mallocz(sizeof(*tq));
  86     if (!tq)
  87         return NULL;
  88
  89     ret = pthread_cond_init(&tq->cond, NULL);
  90     if (ret) {
  91         av_freep(&tq);
  92         return NULL;
  93     }
  94
  95     ret = pthread_mutex_init(&tq->lock, NULL);
  96     if (ret) {
  97         pthread_cond_destroy(&tq->cond);
  98         av_freep(&tq);
  99         return NULL;
 100     }
 101
 102     tq->finished = av_calloc(nb_streams, sizeof(*tq->finished));
 103     if (!tq->finished)
 104         goto fail;
 105     tq->nb_streams = nb_streams;
 106
 107     tq->fifo = av_fifo_alloc2(queue_size, sizeof(FifoElem), 0);
 108     if (!tq->fifo)
 109         goto fail;
 110
 111     tq->obj_pool = obj_pool;
 112     tq->obj_move = obj_move;
 113
 114     return tq;
 115 fail:
 116     tq_free(&tq);
 117     return NULL;
 118 }
 119
 120 int tq_send(ThreadQueue *tq, unsigned int stream_idx, void *data)
 121 {
 122     int *finished;
 123     int ret;
 124
 125     av_assert0(stream_idx < tq->nb_streams);
 126     finished = &tq->finished[stream_idx];
 127
 128     pthread_mutex_lock(&tq->lock);
 129
 130     if (*finished & FINISHED_SEND) {
 131         ret = AVERROR(EINVAL);
 132         goto finish;
 133     }
 134
 135     while (!(*finished & FINISHED_RECV) && !av_fifo_can_write(tq->fifo))
 136         pthread_cond_wait(&tq->cond, &tq->lock);
 137
 138     if (*finished & FINISHED_RECV) {
 139         ret = AVERROR_EOF;
 140         *finished |= FINISHED_SEND;
 141     } else {
 142         FifoElem elem = { .stream_idx = stream_idx };
 143
 144         ret = objpool_get(tq->obj_pool, &elem.obj);
 145         if (ret < 0)
 146             goto finish;
 147
 148         tq->obj_move(elem.obj, data);
 149
 150         ret = av_fifo_write(tq->fifo, &elem, 1);
 151         av_assert0(ret >= 0);
 152         pthread_cond_broadcast(&tq->cond);
 153     }
 154
 155 finish:
 156     pthread_mutex_unlock(&tq->lock);
 157
 158     return ret;
 159 }
 160
 161 static int receive_locked(ThreadQueue *tq, int *stream_idx,
 162                           void *data)
 163 {
 164     FifoElem elem;
 165     unsigned int nb_finished = 0;
 166
 167     while (av_fifo_read(tq->fifo, &elem, 1) >= 0) {
 168         if (tq->finished[elem.stream_idx] & FINISHED_RECV) {
 169             objpool_release(tq->obj_pool, &elem.obj);
 170             continue;
 171         }
 172
 173         tq->obj_move(data, elem.obj);
 174         objpool_release(tq->obj_pool, &elem.obj);
 175         *stream_idx = elem.stream_idx;
 176         return 0;
 177     }
 178
 179     for (unsigned int i = 0; i < tq->nb_streams; i++) {
 180         if (!tq->finished[i])
 181             continue;
 182
 183         /* return EOF to the consumer at most once for each stream */
 184         if (!(tq->finished[i] & FINISHED_RECV)) {
 185             tq->finished[i] |= FINISHED_RECV;
 186             *stream_idx   = i;
 187             return AVERROR_EOF;
 188         }
 189
 190         nb_finished++;
 191     }
 192
 193     return nb_finished == tq->nb_streams ? AVERROR_EOF : AVERROR(EAGAIN);
 194 }
 195
 196 int tq_receive(ThreadQueue *tq, int *stream_idx, void *data)
 197 {
 198     int ret;
 199
 200     *stream_idx = -1;
 201
 202     pthread_mutex_lock(&tq->lock);
 203
 204     while (1) {
 205         size_t can_read = av_fifo_can_read(tq->fifo);
 206
 207         ret = receive_locked(tq, stream_idx, data);
 208
 209         // signal other threads if the fifo state changed
 210         if (can_read != av_fifo_can_read(tq->fifo))
 211             pthread_cond_broadcast(&tq->cond);
 212
 213         if (ret == AVERROR(EAGAIN)) {
 214             pthread_cond_wait(&tq->cond, &tq->lock);
 215             continue;
 216         }
 217
 218         break;
 219     }
 220
 221     pthread_mutex_unlock(&tq->lock);
 222
 223     return ret;
 224 }
 225
 226 void tq_send_finish(ThreadQueue *tq, unsigned int stream_idx)
 227 {
 228     av_assert0(stream_idx < tq->nb_streams);
 229
 230     pthread_mutex_lock(&tq->lock);
 231
 232     /* mark the stream as send-finished;
 233      * next time the consumer thread tries to read this stream it will get
 234      * an EOF and recv-finished flag will be set */
 235     tq->finished[stream_idx] |= FINISHED_SEND;
 236     pthread_cond_broadcast(&tq->cond);
 237
 238     pthread_mutex_unlock(&tq->lock);
 239 }
 240
 241 void tq_receive_finish(ThreadQueue *tq, unsigned int stream_idx)
 242 {
 243     av_assert0(stream_idx < tq->nb_streams);
 244
 245     pthread_mutex_lock(&tq->lock);
 246
 247     /* mark the stream as recv-finished;
 248      * next time the producer thread tries to send for this stream, it will
 249      * get an EOF and send-finished flag will be set */
 250     tq->finished[stream_idx] |= FINISHED_RECV;
 251     pthread_cond_broadcast(&tq->cond);
 252
 253     pthread_mutex_unlock(&tq->lock);
 254 }