avio: Allow custom IO users to get labels for the output bytestream

This allows callers with avio write callbacks to get the bytestream
positions that correspond to keyframes, suitable for live streaming.

In the simplest form, a caller could expect that a header is written
to the bytestream during the avformat_write_header, and the data
output to the avio context during e.g. av_write_frame corresponds
exactly to the current packet passed in.

When combined with av_interleaved_write_frame, and with muxers that
do buffering (most muxers that do some sort of fragmenting or
clustering), the mapping from input data to bytestream positions
is nontrivial.

This allows callers to get directly information about what part
of the bytestream is what, without having to resort to assumptions
about the muxer behaviour.

One keyframe/fragment/block can still be split into multiple (if
they are larger than the aviocontext buffer), which would call
the callback with e.g. AVIO_DATA_MARKER_SYNC_POINT, followed by
AVIO_DATA_MARKER_UNKNOWN for the second time it is called with
the following data.

Signed-off-by: Martin Storsjö <martin@martin.st>

1 files changed, 50 insertions, 2 deletions

diff --git a/libavformat/aviobuf.c b/libavformat/aviobuf.c
index 29fccbea83..706cf5dd7f 100644
--- a/libavformat/aviobuf.c
+++ b/libavformat/aviobuf.c
@@ -140,6 +140,11 @@ int ffio_init_context(AVIOContext *s,
     s->read_pause = NULL;
     s->read_seek  = NULL;
 
+    s->write_data_type       = NULL;
+    s->ignore_boundary_point = 0;
+    s->current_type          = AVIO_DATA_MARKER_UNKNOWN;
+    s->last_time             = AV_NOPTS_VALUE;
+
     return 0;
 }
 
@@ -163,13 +168,25 @@ AVIOContext *avio_alloc_context(
 static void flush_buffer(AVIOContext *s)
 {
     if (s->buf_ptr > s->buffer) {
-        if (s->write_packet && !s->error) {
-            int ret = s->write_packet(s->opaque, s->buffer,
+        if (!s->error) {
+            int ret = 0;
+            if (s->write_data_type)
+                ret = s->write_data_type(s->opaque, s->buffer,
+                                         s->buf_ptr - s->buffer,
+                                         s->current_type,
+                                         s->last_time);
+            else if (s->write_packet)
+                ret = s->write_packet(s->opaque, s->buffer,
                                       s->buf_ptr - s->buffer);
             if (ret < 0) {
                 s->error = ret;
             }
         }
+        if (s->current_type == AVIO_DATA_MARKER_SYNC_POINT ||
+            s->current_type == AVIO_DATA_MARKER_BOUNDARY_POINT) {
+            s->current_type = AVIO_DATA_MARKER_UNKNOWN;
+        }
+        s->last_time = AV_NOPTS_VALUE;
         if (s->update_checksum) {
             s->checksum     = s->update_checksum(s->checksum, s->checksum_ptr,
                                                  s->buf_ptr - s->checksum_ptr);
@@ -402,6 +419,37 @@ void avio_wb24(AVIOContext *s, unsigned int val)
     avio_w8(s, val);
 }
 
+void avio_write_marker(AVIOContext *s, int64_t time, enum AVIODataMarkerType type)
+{
+    if (!s->write_data_type)
+        return;
+    // If ignoring boundary points, just treat it as unknown
+    if (type == AVIO_DATA_MARKER_BOUNDARY_POINT && s->ignore_boundary_point)
+        type = AVIO_DATA_MARKER_UNKNOWN;
+    // Avoid unnecessary flushes if we are already in non-header/trailer
+    // data and setting the type to unknown
+    if (type == AVIO_DATA_MARKER_UNKNOWN &&
+        (s->current_type != AVIO_DATA_MARKER_HEADER &&
+         s->current_type != AVIO_DATA_MARKER_TRAILER))
+        return;
+
+    switch (type) {
+    case AVIO_DATA_MARKER_HEADER:
+    case AVIO_DATA_MARKER_TRAILER:
+        // For header/trailer, ignore a new marker of the same type;
+        // consecutive header/trailer markers can be merged.
+        if (type == s->current_type)
+            return;
+        break;
+    }
+
+    // If we've reached here, we have a new, noteworthy marker.
+    // Flush the previous data and mark the start of the new data.
+    avio_flush(s);
+    s->current_type = type;
+    s->last_time = time;
+}
+
 /* Input stream */
 
 static void fill_buffer(AVIOContext *s)