avformat/tee: Support arbitrary number of slaves
[ffmpeg.git] / libavformat / tee.c
1 /*
2  * Tee pseudo-muxer
3  * Copyright (c) 2012 Nicolas George
4  *
5  * This file is part of FFmpeg.
6  *
7  * FFmpeg is free software; you can redistribute it and/or
8  * modify it under the terms of the GNU Lesser General Public License
9  * as published by the Free Software Foundation; either
10  * version 2.1 of the License, or (at your option) any later version.
11  *
12  * FFmpeg is distributed in the hope that it will be useful,
13  * but WITHOUT ANY WARRANTY; without even the implied warranty of
14  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
15  * GNU Lesser General Public License for more details.
16  *
17  * You should have received a copy of the GNU Lesser General Public License
18  * along with FFmpeg; if not, write to the Free Software * Foundation, Inc.,
19  * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
20  */
21
22
23 #include "libavutil/avutil.h"
24 #include "libavutil/avstring.h"
25 #include "libavutil/opt.h"
26 #include "internal.h"
27 #include "avformat.h"
28 #include "avio_internal.h"
29
30 typedef enum {
31     ON_SLAVE_FAILURE_ABORT  = 1,
32     ON_SLAVE_FAILURE_IGNORE = 2
33 } SlaveFailurePolicy;
34
35 #define DEFAULT_SLAVE_FAILURE_POLICY ON_SLAVE_FAILURE_ABORT
36
37 typedef struct {
38     AVFormatContext *avf;
39     AVBitStreamFilterContext **bsfs; ///< bitstream filters per stream
40
41     SlaveFailurePolicy on_fail;
42
43     /** map from input to output streams indexes,
44      * disabled output streams are set to -1 */
45     int *stream_map;
46     int header_written;
47 } TeeSlave;
48
49 typedef struct TeeContext {
50     const AVClass *class;
51     unsigned nb_slaves;
52     unsigned nb_alive;
53     TeeSlave *slaves;
54 } TeeContext;
55
56 static const char *const slave_delim     = "|";
57 static const char *const slave_opt_open  = "[";
58 static const char *const slave_opt_close = "]";
59 static const char *const slave_opt_delim = ":]"; /* must have the close too */
60 static const char *const slave_bsfs_spec_sep = "/";
61 static const char *const slave_select_sep = ",";
62
63 static const AVClass tee_muxer_class = {
64     .class_name = "Tee muxer",
65     .item_name  = av_default_item_name,
66     .version    = LIBAVUTIL_VERSION_INT,
67 };
68
69 static int parse_slave_options(void *log, char *slave,
70                                AVDictionary **options, char **filename)
71 {
72     const char *p;
73     char *key, *val;
74     int ret;
75
76     if (!strspn(slave, slave_opt_open)) {
77         *filename = slave;
78         return 0;
79     }
80     p = slave + 1;
81     if (strspn(p, slave_opt_close)) {
82         *filename = (char *)p + 1;
83         return 0;
84     }
85     while (1) {
86         ret = av_opt_get_key_value(&p, "=", slave_opt_delim, 0, &key, &val);
87         if (ret < 0) {
88             av_log(log, AV_LOG_ERROR, "No option found near \"%s\"\n", p);
89             goto fail;
90         }
91         ret = av_dict_set(options, key, val,
92                           AV_DICT_DONT_STRDUP_KEY | AV_DICT_DONT_STRDUP_VAL);
93         if (ret < 0)
94             goto fail;
95         if (strspn(p, slave_opt_close))
96             break;
97         p++;
98     }
99     *filename = (char *)p + 1;
100     return 0;
101
102 fail:
103     av_dict_free(options);
104     return ret;
105 }
106
107 /**
108  * Parse list of bitstream filters and add them to the list of filters
109  * pointed to by bsfs.
110  *
111  * The list must be specified in the form:
112  * BSFS ::= BSF[,BSFS]
113  */
114 static int parse_bsfs(void *log_ctx, const char *bsfs_spec,
115                       AVBitStreamFilterContext **bsfs)
116 {
117     char *bsf_name, *buf, *dup, *saveptr;
118     int ret = 0;
119
120     if (!(dup = buf = av_strdup(bsfs_spec)))
121         return AVERROR(ENOMEM);
122
123     while (bsf_name = av_strtok(buf, ",", &saveptr)) {
124         AVBitStreamFilterContext *bsf = av_bitstream_filter_init(bsf_name);
125
126         if (!bsf) {
127             av_log(log_ctx, AV_LOG_ERROR,
128                    "Cannot initialize bitstream filter with name '%s', "
129                    "unknown filter or internal error happened\n",
130                    bsf_name);
131             ret = AVERROR_UNKNOWN;
132             goto end;
133         }
134
135         /* append bsf context to the list of bsf contexts */
136         *bsfs = bsf;
137         bsfs = &bsf->next;
138
139         buf = NULL;
140     }
141
142 end:
143     av_free(dup);
144     return ret;
145 }
146
147 static inline int parse_slave_failure_policy_option(const char *opt, TeeSlave *tee_slave)
148 {
149     if (!opt) {
150         tee_slave->on_fail = DEFAULT_SLAVE_FAILURE_POLICY;
151         return 0;
152     } else if (!av_strcasecmp("abort", opt)) {
153         tee_slave->on_fail = ON_SLAVE_FAILURE_ABORT;
154         return 0;
155     } else if (!av_strcasecmp("ignore", opt)) {
156         tee_slave->on_fail = ON_SLAVE_FAILURE_IGNORE;
157         return 0;
158     }
159     /* Set failure behaviour to abort, so invalid option error will not be ignored */
160     tee_slave->on_fail = ON_SLAVE_FAILURE_ABORT;
161     return AVERROR(EINVAL);
162 }
163
164 static int close_slave(TeeSlave *tee_slave)
165 {
166     AVFormatContext *avf;
167     unsigned i;
168     int ret = 0;
169
170     avf = tee_slave->avf;
171     if (!avf)
172         return 0;
173
174     if (tee_slave->header_written)
175         ret = av_write_trailer(avf);
176
177     if (tee_slave->bsfs) {
178         for (i = 0; i < avf->nb_streams; ++i) {
179             AVBitStreamFilterContext *bsf_next, *bsf = tee_slave->bsfs[i];
180             while (bsf) {
181                 bsf_next = bsf->next;
182                 av_bitstream_filter_close(bsf);
183                 bsf = bsf_next;
184             }
185         }
186     }
187     av_freep(&tee_slave->stream_map);
188     av_freep(&tee_slave->bsfs);
189
190     ff_format_io_close(avf, &avf->pb);
191     avformat_free_context(avf);
192     tee_slave->avf = NULL;
193     return ret;
194 }
195
196 static void close_slaves(AVFormatContext *avf)
197 {
198     TeeContext *tee = avf->priv_data;
199     unsigned i;
200
201     for (i = 0; i < tee->nb_slaves; i++) {
202         close_slave(&tee->slaves[i]);
203     }
204     av_freep(&tee->slaves);
205 }
206
207 static int open_slave(AVFormatContext *avf, char *slave, TeeSlave *tee_slave)
208 {
209     int i, ret;
210     AVDictionary *options = NULL;
211     AVDictionaryEntry *entry;
212     char *filename;
213     char *format = NULL, *select = NULL, *on_fail = NULL;
214     AVFormatContext *avf2 = NULL;
215     AVStream *st, *st2;
216     int stream_count;
217     int fullret;
218     char *subselect = NULL, *next_subselect = NULL, *first_subselect = NULL, *tmp_select = NULL;
219
220     if ((ret = parse_slave_options(avf, slave, &options, &filename)) < 0)
221         return ret;
222
223 #define STEAL_OPTION(option, field) do {                                \
224         if ((entry = av_dict_get(options, option, NULL, 0))) {          \
225             field = entry->value;                                       \
226             entry->value = NULL; /* prevent it from being freed */      \
227             av_dict_set(&options, option, NULL, 0);                     \
228         }                                                               \
229     } while (0)
230
231     STEAL_OPTION("f", format);
232     STEAL_OPTION("select", select);
233     STEAL_OPTION("onfail", on_fail);
234
235     ret = parse_slave_failure_policy_option(on_fail, tee_slave);
236     if (ret < 0) {
237         av_log(avf, AV_LOG_ERROR,
238                "Invalid onfail option value, valid options are 'abort' and 'ignore'\n");
239         goto end;
240     }
241
242     ret = avformat_alloc_output_context2(&avf2, NULL, format, filename);
243     if (ret < 0)
244         goto end;
245     tee_slave->avf = avf2;
246     av_dict_copy(&avf2->metadata, avf->metadata, 0);
247     avf2->opaque   = avf->opaque;
248     avf2->io_open  = avf->io_open;
249     avf2->io_close = avf->io_close;
250
251     tee_slave->stream_map = av_calloc(avf->nb_streams, sizeof(*tee_slave->stream_map));
252     if (!tee_slave->stream_map) {
253         ret = AVERROR(ENOMEM);
254         goto end;
255     }
256
257     stream_count = 0;
258     for (i = 0; i < avf->nb_streams; i++) {
259         st = avf->streams[i];
260         if (select) {
261             tmp_select = av_strdup(select);  // av_strtok is destructive so we regenerate it in each loop
262             if (!tmp_select) {
263                 ret = AVERROR(ENOMEM);
264                 goto end;
265             }
266             fullret = 0;
267             first_subselect = tmp_select;
268             next_subselect = NULL;
269             while (subselect = av_strtok(first_subselect, slave_select_sep, &next_subselect)) {
270                 first_subselect = NULL;
271
272                 ret = avformat_match_stream_specifier(avf, avf->streams[i], subselect);
273                 if (ret < 0) {
274                     av_log(avf, AV_LOG_ERROR,
275                            "Invalid stream specifier '%s' for output '%s'\n",
276                            subselect, slave);
277                     goto end;
278                 }
279                 if (ret != 0) {
280                     fullret = 1; // match
281                     break;
282                 }
283             }
284             av_freep(&tmp_select);
285
286             if (fullret == 0) { /* no match */
287                 tee_slave->stream_map[i] = -1;
288                 continue;
289             }
290         }
291         tee_slave->stream_map[i] = stream_count++;
292
293         if (!(st2 = avformat_new_stream(avf2, NULL))) {
294             ret = AVERROR(ENOMEM);
295             goto end;
296         }
297         st2->id = st->id;
298         st2->r_frame_rate        = st->r_frame_rate;
299         st2->time_base           = st->time_base;
300         st2->start_time          = st->start_time;
301         st2->duration            = st->duration;
302         st2->nb_frames           = st->nb_frames;
303         st2->disposition         = st->disposition;
304         st2->sample_aspect_ratio = st->sample_aspect_ratio;
305         st2->avg_frame_rate      = st->avg_frame_rate;
306         av_dict_copy(&st2->metadata, st->metadata, 0);
307         if ((ret = avcodec_parameters_copy(st2->codecpar, st->codecpar)) < 0)
308             goto end;
309     }
310
311     if (!(avf2->oformat->flags & AVFMT_NOFILE)) {
312         if ((ret = avf2->io_open(avf2, &avf2->pb, filename, AVIO_FLAG_WRITE, NULL)) < 0) {
313             av_log(avf, AV_LOG_ERROR, "Slave '%s': error opening: %s\n",
314                    slave, av_err2str(ret));
315             goto end;
316         }
317     }
318
319     if ((ret = avformat_write_header(avf2, &options)) < 0) {
320         av_log(avf, AV_LOG_ERROR, "Slave '%s': error writing header: %s\n",
321                slave, av_err2str(ret));
322         goto end;
323     }
324     tee_slave->header_written = 1;
325
326     tee_slave->bsfs = av_calloc(avf2->nb_streams, sizeof(*tee_slave->bsfs));
327     if (!tee_slave->bsfs) {
328         ret = AVERROR(ENOMEM);
329         goto end;
330     }
331
332     entry = NULL;
333     while (entry = av_dict_get(options, "bsfs", NULL, AV_DICT_IGNORE_SUFFIX)) {
334         const char *spec = entry->key + strlen("bsfs");
335         if (*spec) {
336             if (strspn(spec, slave_bsfs_spec_sep) != 1) {
337                 av_log(avf, AV_LOG_ERROR,
338                        "Specifier separator in '%s' is '%c', but only characters '%s' "
339                        "are allowed\n", entry->key, *spec, slave_bsfs_spec_sep);
340                 ret = AVERROR(EINVAL);
341                 goto end;
342             }
343             spec++; /* consume separator */
344         }
345
346         for (i = 0; i < avf2->nb_streams; i++) {
347             ret = avformat_match_stream_specifier(avf2, avf2->streams[i], spec);
348             if (ret < 0) {
349                 av_log(avf, AV_LOG_ERROR,
350                        "Invalid stream specifier '%s' in bsfs option '%s' for slave "
351                        "output '%s'\n", spec, entry->key, filename);
352                 goto end;
353             }
354
355             if (ret > 0) {
356                 av_log(avf, AV_LOG_DEBUG, "spec:%s bsfs:%s matches stream %d of slave "
357                        "output '%s'\n", spec, entry->value, i, filename);
358                 if (tee_slave->bsfs[i]) {
359                     av_log(avf, AV_LOG_WARNING,
360                            "Duplicate bsfs specification associated to stream %d of slave "
361                            "output '%s', filters will be ignored\n", i, filename);
362                     continue;
363                 }
364                 ret = parse_bsfs(avf, entry->value, &tee_slave->bsfs[i]);
365                 if (ret < 0) {
366                     av_log(avf, AV_LOG_ERROR,
367                            "Error parsing bitstream filter sequence '%s' associated to "
368                            "stream %d of slave output '%s'\n", entry->value, i, filename);
369                     goto end;
370                 }
371             }
372         }
373
374         av_dict_set(&options, entry->key, NULL, 0);
375     }
376
377     if (options) {
378         entry = NULL;
379         while ((entry = av_dict_get(options, "", entry, AV_DICT_IGNORE_SUFFIX)))
380             av_log(avf2, AV_LOG_ERROR, "Unknown option '%s'\n", entry->key);
381         ret = AVERROR_OPTION_NOT_FOUND;
382         goto end;
383     }
384
385 end:
386     av_free(format);
387     av_free(select);
388     av_free(on_fail);
389     av_dict_free(&options);
390     av_freep(&tmp_select);
391     return ret;
392 }
393
394 static void log_slave(TeeSlave *slave, void *log_ctx, int log_level)
395 {
396     int i;
397     av_log(log_ctx, log_level, "filename:'%s' format:%s\n",
398            slave->avf->filename, slave->avf->oformat->name);
399     for (i = 0; i < slave->avf->nb_streams; i++) {
400         AVStream *st = slave->avf->streams[i];
401         AVBitStreamFilterContext *bsf = slave->bsfs[i];
402
403         av_log(log_ctx, log_level, "    stream:%d codec:%s type:%s",
404                i, avcodec_get_name(st->codecpar->codec_id),
405                av_get_media_type_string(st->codecpar->codec_type));
406         if (bsf) {
407             av_log(log_ctx, log_level, " bsfs:");
408             while (bsf) {
409                 av_log(log_ctx, log_level, "%s%s",
410                        bsf->filter->name, bsf->next ? "," : "");
411                 bsf = bsf->next;
412             }
413         }
414         av_log(log_ctx, log_level, "\n");
415     }
416 }
417
418 static int tee_process_slave_failure(AVFormatContext *avf, unsigned slave_idx, int err_n)
419 {
420     TeeContext *tee = avf->priv_data;
421     TeeSlave *tee_slave = &tee->slaves[slave_idx];
422
423     tee->nb_alive--;
424
425     close_slave(tee_slave);
426
427     if (!tee->nb_alive) {
428         av_log(avf, AV_LOG_ERROR, "All tee outputs failed.\n");
429         return err_n;
430     } else if (tee_slave->on_fail == ON_SLAVE_FAILURE_ABORT) {
431         av_log(avf, AV_LOG_ERROR, "Slave muxer #%u failed, aborting.\n", slave_idx);
432         return err_n;
433     } else {
434         av_log(avf, AV_LOG_ERROR, "Slave muxer #%u failed: %s, continuing with %u/%u slaves.\n",
435                slave_idx, av_err2str(err_n), tee->nb_alive, tee->nb_slaves);
436         return 0;
437     }
438 }
439
440 static int tee_write_header(AVFormatContext *avf)
441 {
442     TeeContext *tee = avf->priv_data;
443     unsigned nb_slaves = 0, i;
444     const char *filename = avf->filename;
445     char **slaves = NULL;
446     int ret;
447
448     while (*filename) {
449         char *slave = av_get_token(&filename, slave_delim);
450         if (!slave) {
451             ret = AVERROR(ENOMEM);
452             goto fail;
453         }
454         ret = av_dynarray_add_nofree(&slaves, &nb_slaves, slave);
455         if (ret < 0) {
456             av_free(slave);
457             goto fail;
458         }
459         if (strspn(filename, slave_delim))
460             filename++;
461     }
462
463     if (!(tee->slaves = av_mallocz_array(nb_slaves, sizeof(*tee->slaves)))) {
464         ret = AVERROR(ENOMEM);
465         goto fail;
466     }
467     tee->nb_slaves = tee->nb_alive = nb_slaves;
468
469     for (i = 0; i < nb_slaves; i++) {
470         if ((ret = open_slave(avf, slaves[i], &tee->slaves[i])) < 0) {
471             ret = tee_process_slave_failure(avf, i, ret);
472             if (ret < 0)
473                 goto fail;
474         } else {
475             log_slave(&tee->slaves[i], avf, AV_LOG_VERBOSE);
476         }
477         av_freep(&slaves[i]);
478     }
479
480     for (i = 0; i < avf->nb_streams; i++) {
481         int j, mapped = 0;
482         for (j = 0; j < tee->nb_slaves; j++)
483             if (tee->slaves[j].avf)
484                 mapped += tee->slaves[j].stream_map[i] >= 0;
485         if (!mapped)
486             av_log(avf, AV_LOG_WARNING, "Input stream #%d is not mapped "
487                    "to any slave.\n", i);
488     }
489     av_free(slaves);
490     return 0;
491
492 fail:
493     for (i = 0; i < nb_slaves; i++)
494         av_freep(&slaves[i]);
495     close_slaves(avf);
496     av_free(slaves);
497     return ret;
498 }
499
500 static int tee_write_trailer(AVFormatContext *avf)
501 {
502     TeeContext *tee = avf->priv_data;
503     int ret_all = 0, ret;
504     unsigned i;
505
506     for (i = 0; i < tee->nb_slaves; i++) {
507         if ((ret = close_slave(&tee->slaves[i])) < 0) {
508             ret = tee_process_slave_failure(avf, i, ret);
509             if (!ret_all && ret < 0)
510                 ret_all = ret;
511         }
512     }
513     av_freep(&tee->slaves);
514     return ret_all;
515 }
516
517 static int tee_write_packet(AVFormatContext *avf, AVPacket *pkt)
518 {
519     TeeContext *tee = avf->priv_data;
520     AVFormatContext *avf2;
521     AVPacket pkt2;
522     int ret_all = 0, ret;
523     unsigned i, s;
524     int s2;
525     AVRational tb, tb2;
526
527     for (i = 0; i < tee->nb_slaves; i++) {
528         if (!(avf2 = tee->slaves[i].avf))
529             continue;
530
531         s = pkt->stream_index;
532         s2 = tee->slaves[i].stream_map[s];
533         if (s2 < 0)
534             continue;
535
536         memset(&pkt2, 0, sizeof(AVPacket));
537         if ((ret = av_packet_ref(&pkt2, pkt)) < 0)
538             if (!ret_all) {
539                 ret_all = ret;
540                 continue;
541             }
542         tb  = avf ->streams[s ]->time_base;
543         tb2 = avf2->streams[s2]->time_base;
544         pkt2.pts      = av_rescale_q(pkt->pts,      tb, tb2);
545         pkt2.dts      = av_rescale_q(pkt->dts,      tb, tb2);
546         pkt2.duration = av_rescale_q(pkt->duration, tb, tb2);
547         pkt2.stream_index = s2;
548
549         if ((ret = av_apply_bitstream_filters(avf2->streams[s2]->codec, &pkt2,
550                                               tee->slaves[i].bsfs[s2])) < 0 ||
551             (ret = av_interleaved_write_frame(avf2, &pkt2)) < 0) {
552             ret = tee_process_slave_failure(avf, i, ret);
553             if (!ret_all && ret < 0)
554                 ret_all = ret;
555         }
556     }
557     return ret_all;
558 }
559
560 AVOutputFormat ff_tee_muxer = {
561     .name              = "tee",
562     .long_name         = NULL_IF_CONFIG_SMALL("Multiple muxer tee"),
563     .priv_data_size    = sizeof(TeeContext),
564     .write_header      = tee_write_header,
565     .write_trailer     = tee_write_trailer,
566     .write_packet      = tee_write_packet,
567     .priv_class        = &tee_muxer_class,
568     .flags             = AVFMT_NOFILE,
569 };