avcodec/cbs: Factor out common code for writing units

All cbs-functions to write units share a common pattern:
1. They check whether they have a write buffer (that is used to store
the unit's data until the needed size becomes known after writing the
unit when a dedicated buffer will be allocated).
2. They use this buffer for a PutBitContext.
3. The (codec-specific) writing takes place through the PutBitContext.
4. The return value is checked. AVERROR(ENOSPC) here always indicates
that the buffer was too small and leads to a reallocation of said
buffer.
5. The final buffer will be allocated and the data copied.

This commit factors this common code out in a single function in cbs.c.

Signed-off-by: Andreas Rheinhardt <andreas.rheinhardt@gmail.com>
(cherry picked from commit 7c92eaace2b338e0b3acc18e1543b365610578fd)

13 files changed, 94 insertions, 322 deletions

diff --git a/libavcodec/cbs.c b/libavcodec/cbs.c
index 2350416501..047f4fb898 100644
--- a/libavcodec/cbs.c
+++ b/libavcodec/cbs.c
@@ -95,10 +95,12 @@ int ff_cbs_init(CodedBitstreamContext **ctx_ptr,
     ctx->log_ctx = log_ctx;
     ctx->codec   = type;
 
-    ctx->priv_data = av_mallocz(ctx->codec->priv_data_size);
-    if (!ctx->priv_data) {
-        av_freep(&ctx);
-        return AVERROR(ENOMEM);
+    if (type->priv_data_size) {
+        ctx->priv_data = av_mallocz(ctx->codec->priv_data_size);
+        if (!ctx->priv_data) {
+            av_freep(&ctx);
+            return AVERROR(ENOMEM);
+        }
     }
 
     ctx->decompose_unit_types = NULL;
@@ -120,6 +122,7 @@ void ff_cbs_close(CodedBitstreamContext **ctx_ptr)
     if (ctx->codec && ctx->codec->close)
         ctx->codec->close(ctx);
 
+    av_freep(&ctx->write_buffer);
     av_freep(&ctx->priv_data);
     av_freep(ctx_ptr);
 }
@@ -280,6 +283,57 @@ int ff_cbs_read(CodedBitstreamContext *ctx,
     return cbs_read_fragment_content(ctx, frag);
 }
 
+static int cbs_write_unit_data(CodedBitstreamContext *ctx,
+                               CodedBitstreamUnit *unit)
+{
+    PutBitContext pbc;
+    int ret;
+
+    if (!ctx->write_buffer) {
+        // Initial write buffer size is 1MB.
+        ctx->write_buffer_size = 1024 * 1024;
+
+    reallocate_and_try_again:
+        ret = av_reallocp(&ctx->write_buffer, ctx->write_buffer_size);
+        if (ret < 0) {
+            av_log(ctx->log_ctx, AV_LOG_ERROR, "Unable to allocate a "
+                   "sufficiently large write buffer (last attempt "
+                   "%"SIZE_SPECIFIER" bytes).\n", ctx->write_buffer_size);
+            return ret;
+        }
+    }
+
+    init_put_bits(&pbc, ctx->write_buffer, ctx->write_buffer_size);
+
+    ret = ctx->codec->write_unit(ctx, unit, &pbc);
+    if (ret < 0) {
+        if (ret == AVERROR(ENOSPC)) {
+            // Overflow.
+            ctx->write_buffer_size *= 2;
+            goto reallocate_and_try_again;
+        }
+        // Write failed for some other reason.
+        return ret;
+    }
+
+    // Overflow but we didn't notice.
+    av_assert0(put_bits_count(&pbc) <= 8 * ctx->write_buffer_size);
+
+    if (put_bits_count(&pbc) % 8)
+        unit->data_bit_padding = 8 - put_bits_count(&pbc) % 8;
+    else
+        unit->data_bit_padding = 0;
+
+    flush_put_bits(&pbc);
+
+    ret = ff_cbs_alloc_unit_data(ctx, unit, put_bits_count(&pbc) / 8);
+    if (ret < 0)
+        return ret;
+
+    memcpy(unit->data, ctx->write_buffer, unit->data_size);
+
+    return 0;
+}
 
 int ff_cbs_write_fragment_data(CodedBitstreamContext *ctx,
                                CodedBitstreamFragment *frag)
@@ -295,7 +349,7 @@ int ff_cbs_write_fragment_data(CodedBitstreamContext *ctx,
         av_buffer_unref(&unit->data_ref);
         unit->data = NULL;
 
-        err = ctx->codec->write_unit(ctx, unit);
+        err = cbs_write_unit_data(ctx, unit);
         if (err < 0) {
             av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to write unit %d "
                    "(type %"PRIu32").\n", i, unit->type);
diff --git a/libavcodec/cbs.h b/libavcodec/cbs.h
index fe57e7b2a5..1974cf3e66 100644
--- a/libavcodec/cbs.h
+++ b/libavcodec/cbs.h
@@ -210,6 +210,13 @@ typedef struct CodedBitstreamContext {
      * From AV_LOG_*; defaults to AV_LOG_TRACE.
      */
     int trace_level;
+
+    /**
+     * Write buffer. Used as intermediate buffer when writing units.
+     * For internal use of cbs only.
+     */
+    uint8_t *write_buffer;
+    size_t   write_buffer_size;
 } CodedBitstreamContext;
 
 
diff --git a/libavcodec/cbs_av1.c b/libavcodec/cbs_av1.c
index cf3561a4ea..ffb68adcd3 100644
--- a/libavcodec/cbs_av1.c
+++ b/libavcodec/cbs_av1.c
@@ -1205,66 +1205,19 @@ static int cbs_av1_write_obu(CodedBitstreamContext *ctx,
         return AVERROR(ENOSPC);
 
     if (obu->obu_size > 0) {
-        memmove(priv->write_buffer + data_pos,
-                priv->write_buffer + start_pos, header_size);
+        memmove(pbc->buf + data_pos,
+                pbc->buf + start_pos, header_size);
         skip_put_bytes(pbc, header_size);
 
         if (td) {
-            memcpy(priv->write_buffer + data_pos + header_size,
+            memcpy(pbc->buf + data_pos + header_size,
                    td->data, td->data_size);
             skip_put_bytes(pbc, td->data_size);
         }
     }
 
-    return 0;
-}
-
-static int cbs_av1_write_unit(CodedBitstreamContext *ctx,
-                              CodedBitstreamUnit *unit)
-{
-    CodedBitstreamAV1Context *priv = ctx->priv_data;
-    PutBitContext pbc;
-    int err;
-
-    if (!priv->write_buffer) {
-        // Initial write buffer size is 1MB.
-        priv->write_buffer_size = 1024 * 1024;
-
-    reallocate_and_try_again:
-        err = av_reallocp(&priv->write_buffer, priv->write_buffer_size);
-        if (err < 0) {
-            av_log(ctx->log_ctx, AV_LOG_ERROR, "Unable to allocate a "
-                   "sufficiently large write buffer (last attempt "
-                   "%"SIZE_SPECIFIER" bytes).\n", priv->write_buffer_size);
-            return err;
-        }
-    }
-
-    init_put_bits(&pbc, priv->write_buffer, priv->write_buffer_size);
-
-    err = cbs_av1_write_obu(ctx, unit, &pbc);
-    if (err == AVERROR(ENOSPC)) {
-        // Overflow.
-        priv->write_buffer_size *= 2;
-        goto reallocate_and_try_again;
-    }
-    if (err < 0)
-        return err;
-
-    // Overflow but we didn't notice.
-    av_assert0(put_bits_count(&pbc) <= 8 * priv->write_buffer_size);
-
     // OBU data must be byte-aligned.
-    av_assert0(put_bits_count(&pbc) % 8 == 0);
-
-    unit->data_size = put_bits_count(&pbc) / 8;
-    flush_put_bits(&pbc);
-
-    err = ff_cbs_alloc_unit_data(ctx, unit, unit->data_size);
-    if (err < 0)
-        return err;
-
-    memcpy(unit->data, priv->write_buffer, unit->data_size);
+    av_assert0(put_bits_count(pbc) % 8 == 0);
 
     return 0;
 }
@@ -1303,8 +1256,6 @@ static void cbs_av1_close(CodedBitstreamContext *ctx)
 
     av_buffer_unref(&priv->sequence_header_ref);
     av_buffer_unref(&priv->frame_header_ref);
-
-    av_freep(&priv->write_buffer);
 }
 
 const CodedBitstreamType ff_cbs_type_av1 = {
@@ -1314,7 +1265,7 @@ const CodedBitstreamType ff_cbs_type_av1 = {
 
     .split_fragment    = &cbs_av1_split_fragment,
     .read_unit         = &cbs_av1_read_unit,
-    .write_unit        = &cbs_av1_write_unit,
+    .write_unit        = &cbs_av1_write_obu,
     .assemble_fragment = &cbs_av1_assemble_fragment,
 
     .close             = &cbs_av1_close,
diff --git a/libavcodec/cbs_av1.h b/libavcodec/cbs_av1.h
index 656385b2e4..d142882074 100644
--- a/libavcodec/cbs_av1.h
+++ b/libavcodec/cbs_av1.h
@@ -444,10 +444,6 @@ typedef struct CodedBitstreamAV1Context {
     AV1ReferenceFrameState *ref;
     AV1ReferenceFrameState read_ref[AV1_NUM_REF_FRAMES];
     AV1ReferenceFrameState write_ref[AV1_NUM_REF_FRAMES];
-
-    // Write buffer.
-    uint8_t *write_buffer;
-    size_t   write_buffer_size;
 } CodedBitstreamAV1Context;
 
 
diff --git a/libavcodec/cbs_h2645.c b/libavcodec/cbs_h2645.c
index c3cb88d9c6..ba442b62ed 100644
--- a/libavcodec/cbs_h2645.c
+++ b/libavcodec/cbs_h2645.c
@@ -1380,65 +1380,6 @@ static int cbs_h265_write_nal_unit(CodedBitstreamContext *ctx,
     return 0;
 }
 
-static int cbs_h2645_write_nal_unit(CodedBitstreamContext *ctx,
-                                    CodedBitstreamUnit *unit)
-{
-    CodedBitstreamH2645Context *priv = ctx->priv_data;
-    enum AVCodecID codec_id = ctx->codec->codec_id;
-    PutBitContext pbc;
-    int err;
-
-    if (!priv->write_buffer) {
-        // Initial write buffer size is 1MB.
-        priv->write_buffer_size = 1024 * 1024;
-
-    reallocate_and_try_again:
-        err = av_reallocp(&priv->write_buffer, priv->write_buffer_size);
-        if (err < 0) {
-            av_log(ctx->log_ctx, AV_LOG_ERROR, "Unable to allocate a "
-                   "sufficiently large write buffer (last attempt "
-                   "%"SIZE_SPECIFIER" bytes).\n", priv->write_buffer_size);
-            return err;
-        }
-    }
-
-    init_put_bits(&pbc, priv->write_buffer, priv->write_buffer_size);
-
-    if (codec_id == AV_CODEC_ID_H264)
-        err = cbs_h264_write_nal_unit(ctx, unit, &pbc);
-    else
-        err = cbs_h265_write_nal_unit(ctx, unit, &pbc);
-
-    if (err == AVERROR(ENOSPC)) {
-        // Overflow.
-        priv->write_buffer_size *= 2;
-        goto reallocate_and_try_again;
-    }
-    // Overflow but we didn't notice.
-    av_assert0(put_bits_count(&pbc) <= 8 * priv->write_buffer_size);
-
-    if (err < 0) {
-        // Write failed for some other reason.
-        return err;
-    }
-
-    if (put_bits_count(&pbc) % 8)
-        unit->data_bit_padding = 8 - put_bits_count(&pbc) % 8;
-    else
-        unit->data_bit_padding = 0;
-
-    unit->data_size = (put_bits_count(&pbc) + 7) / 8;
-    flush_put_bits(&pbc);
-
-    err = ff_cbs_alloc_unit_data(ctx, unit, unit->data_size);
-    if (err < 0)
-        return err;
-
-    memcpy(unit->data, priv->write_buffer, unit->data_size);
-
-    return 0;
-}
-
 static int cbs_h2645_assemble_fragment(CodedBitstreamContext *ctx,
                                        CodedBitstreamFragment *frag)
 {
@@ -1533,8 +1474,6 @@ static void cbs_h264_close(CodedBitstreamContext *ctx)
 
     ff_h2645_packet_uninit(&h264->common.read_packet);
 
-    av_freep(&h264->common.write_buffer);
-
     for (i = 0; i < FF_ARRAY_ELEMS(h264->sps); i++)
         av_buffer_unref(&h264->sps_ref[i]);
     for (i = 0; i < FF_ARRAY_ELEMS(h264->pps); i++)
@@ -1548,8 +1487,6 @@ static void cbs_h265_close(CodedBitstreamContext *ctx)
 
     ff_h2645_packet_uninit(&h265->common.read_packet);
 
-    av_freep(&h265->common.write_buffer);
-
     for (i = 0; i < FF_ARRAY_ELEMS(h265->vps); i++)
         av_buffer_unref(&h265->vps_ref[i]);
     for (i = 0; i < FF_ARRAY_ELEMS(h265->sps); i++)
@@ -1565,7 +1502,7 @@ const CodedBitstreamType ff_cbs_type_h264 = {
 
     .split_fragment    = &cbs_h2645_split_fragment,
     .read_unit         = &cbs_h264_read_nal_unit,
-    .write_unit        = &cbs_h2645_write_nal_unit,
+    .write_unit        = &cbs_h264_write_nal_unit,
     .assemble_fragment = &cbs_h2645_assemble_fragment,
 
     .close             = &cbs_h264_close,
@@ -1578,7 +1515,7 @@ const CodedBitstreamType ff_cbs_type_h265 = {
 
     .split_fragment    = &cbs_h2645_split_fragment,
     .read_unit         = &cbs_h265_read_nal_unit,
-    .write_unit        = &cbs_h2645_write_nal_unit,
+    .write_unit        = &cbs_h265_write_nal_unit,
     .assemble_fragment = &cbs_h2645_assemble_fragment,
 
     .close             = &cbs_h265_close,
diff --git a/libavcodec/cbs_h2645.h b/libavcodec/cbs_h2645.h
index f4cf65bdde..f4c987a511 100644
--- a/libavcodec/cbs_h2645.h
+++ b/libavcodec/cbs_h2645.h
@@ -19,9 +19,6 @@
 #ifndef AVCODEC_CBS_H2645_H
 #define AVCODEC_CBS_H2645_H
 
-#include <stddef.h>
-#include <stdint.h>
-
 #include "h2645_parse.h"
 
 
@@ -33,10 +30,6 @@ typedef struct CodedBitstreamH2645Context {
     int nal_length_size;
     // Packet reader.
     H2645Packet read_packet;
-
-    // Write buffer
-    uint8_t *write_buffer;
-    size_t write_buffer_size;
 } CodedBitstreamH2645Context;
 
 
diff --git a/libavcodec/cbs_internal.h b/libavcodec/cbs_internal.h
index dd4babf092..4c5a535ca6 100644
--- a/libavcodec/cbs_internal.h
+++ b/libavcodec/cbs_internal.h
@@ -44,9 +44,11 @@ typedef struct CodedBitstreamType {
     int (*read_unit)(CodedBitstreamContext *ctx,
                      CodedBitstreamUnit *unit);
 
-    // Write the unit->data bitstream from unit->content.
+    // Write the data bitstream from unit->content into pbc.
+    // Return value AVERROR(ENOSPC) indicates that pbc was too small.
     int (*write_unit)(CodedBitstreamContext *ctx,
-                      CodedBitstreamUnit *unit);
+                      CodedBitstreamUnit *unit,
+                      PutBitContext *pbc);
 
     // Read the data from all of frag->units and assemble it into
     // a bitstream for the whole fragment.
diff --git a/libavcodec/cbs_jpeg.c b/libavcodec/cbs_jpeg.c
index 83857bbba2..faa8301d4f 100644
--- a/libavcodec/cbs_jpeg.c
+++ b/libavcodec/cbs_jpeg.c
@@ -377,58 +377,13 @@ static int cbs_jpeg_write_segment(CodedBitstreamContext *ctx,
 }
 
 static int cbs_jpeg_write_unit(CodedBitstreamContext *ctx,
-                                CodedBitstreamUnit *unit)
+                               CodedBitstreamUnit *unit,
+                               PutBitContext *pbc)
 {
-    CodedBitstreamJPEGContext *priv = ctx->priv_data;
-    PutBitContext pbc;
-    int err;
-
-    if (!priv->write_buffer) {
-        // Initial write buffer size is 1MB.
-        priv->write_buffer_size = 1024 * 1024;
-
-    reallocate_and_try_again:
-        err = av_reallocp(&priv->write_buffer, priv->write_buffer_size);
-        if (err < 0) {
-            av_log(ctx->log_ctx, AV_LOG_ERROR, "Unable to allocate a "
-                   "sufficiently large write buffer (last attempt "
-                   "%"SIZE_SPECIFIER" bytes).\n", priv->write_buffer_size);
-            return err;
-        }
-    }
-
-    init_put_bits(&pbc, priv->write_buffer, priv->write_buffer_size);
-
     if (unit->type == JPEG_MARKER_SOS)
-        err = cbs_jpeg_write_scan(ctx, unit, &pbc);
-    else
-        err = cbs_jpeg_write_segment(ctx, unit, &pbc);
-
-    if (err == AVERROR(ENOSPC)) {
-        // Overflow.
-        priv->write_buffer_size *= 2;
-        goto reallocate_and_try_again;
-    }
-    if (err < 0) {
-        // Write failed for some other reason.
-        return err;
-    }
-
-    if (put_bits_count(&pbc) % 8)
-        unit->data_bit_padding = 8 - put_bits_count(&pbc) % 8;
+        return cbs_jpeg_write_scan   (ctx, unit, pbc);
     else
-        unit->data_bit_padding = 0;
-
-    unit->data_size = (put_bits_count(&pbc) + 7) / 8;
-    flush_put_bits(&pbc);
-
-    err = ff_cbs_alloc_unit_data(ctx, unit, unit->data_size);
-    if (err < 0)
-        return err;
-
-    memcpy(unit->data, priv->write_buffer, unit->data_size);
-
-    return 0;
+        return cbs_jpeg_write_segment(ctx, unit, pbc);
 }
 
 static int cbs_jpeg_assemble_fragment(CodedBitstreamContext *ctx,
@@ -499,22 +454,11 @@ static int cbs_jpeg_assemble_fragment(CodedBitstreamContext *ctx,
     return 0;
 }
 
-static void cbs_jpeg_close(CodedBitstreamContext *ctx)
-{
-    CodedBitstreamJPEGContext *priv = ctx->priv_data;
-
-    av_freep(&priv->write_buffer);
-}
-
 const CodedBitstreamType ff_cbs_type_jpeg = {
     .codec_id          = AV_CODEC_ID_MJPEG,
 
-    .priv_data_size    = sizeof(CodedBitstreamJPEGContext),
-
     .split_fragment    = &cbs_jpeg_split_fragment,
     .read_unit         = &cbs_jpeg_read_unit,
     .write_unit        = &cbs_jpeg_write_unit,
     .assemble_fragment = &cbs_jpeg_assemble_fragment,
-
-    .close             = &cbs_jpeg_close,
 };
diff --git a/libavcodec/cbs_jpeg.h b/libavcodec/cbs_jpeg.h
index 913d3f90f6..ff1961106f 100644
--- a/libavcodec/cbs_jpeg.h
+++ b/libavcodec/cbs_jpeg.h
@@ -120,11 +120,4 @@ typedef struct JPEGRawComment {
 } JPEGRawComment;
 
 
-typedef struct CodedBitstreamJPEGContext {
-    // Write buffer.
-    uint8_t *write_buffer;
-    size_t write_buffer_size;
-} CodedBitstreamJPEGContext;
-
-
 #endif /* AVCODEC_CBS_JPEG_H */
diff --git a/libavcodec/cbs_mpeg2.c b/libavcodec/cbs_mpeg2.c
index cb202f835b..a49a403b26 100644
--- a/libavcodec/cbs_mpeg2.c
+++ b/libavcodec/cbs_mpeg2.c
@@ -335,58 +335,13 @@ static int cbs_mpeg2_write_slice(CodedBitstreamContext *ctx,
 }
 
 static int cbs_mpeg2_write_unit(CodedBitstreamContext *ctx,
-                                CodedBitstreamUnit *unit)
+                                CodedBitstreamUnit *unit,
+                                PutBitContext *pbc)
 {
-    CodedBitstreamMPEG2Context *priv = ctx->priv_data;
-    PutBitContext pbc;
-    int err;
-
-    if (!priv->write_buffer) {
-        // Initial write buffer size is 1MB.
-        priv->write_buffer_size = 1024 * 1024;
-
-    reallocate_and_try_again:
-        err = av_reallocp(&priv->write_buffer, priv->write_buffer_size);
-        if (err < 0) {
-            av_log(ctx->log_ctx, AV_LOG_ERROR, "Unable to allocate a "
-                   "sufficiently large write buffer (last attempt "
-                   "%"SIZE_SPECIFIER" bytes).\n", priv->write_buffer_size);
-            return err;
-        }
-    }
-
-    init_put_bits(&pbc, priv->write_buffer, priv->write_buffer_size);
-
     if (MPEG2_START_IS_SLICE(unit->type))
-        err = cbs_mpeg2_write_slice(ctx, unit, &pbc);
-    else
-        err = cbs_mpeg2_write_header(ctx, unit, &pbc);
-
-    if (err == AVERROR(ENOSPC)) {
-        // Overflow.
-        priv->write_buffer_size *= 2;
-        goto reallocate_and_try_again;
-    }
-    if (err < 0) {
-        // Write failed for some other reason.
-        return err;
-    }
-
-    if (put_bits_count(&pbc) % 8)
-        unit->data_bit_padding = 8 - put_bits_count(&pbc) % 8;
+        return cbs_mpeg2_write_slice (ctx, unit, pbc);
     else
-        unit->data_bit_padding = 0;
-
-    unit->data_size = (put_bits_count(&pbc) + 7) / 8;
-    flush_put_bits(&pbc);
-
-    err = ff_cbs_alloc_unit_data(ctx, unit, unit->data_size);
-    if (err < 0)
-        return err;
-
-    memcpy(unit->data, priv->write_buffer, unit->data_size);
-
-    return 0;
+        return cbs_mpeg2_write_header(ctx, unit, pbc);
 }
 
 static int cbs_mpeg2_assemble_fragment(CodedBitstreamContext *ctx,
@@ -426,13 +381,6 @@ static int cbs_mpeg2_assemble_fragment(CodedBitstreamContext *ctx,
     return 0;
 }
 
-static void cbs_mpeg2_close(CodedBitstreamContext *ctx)
-{
-    CodedBitstreamMPEG2Context *priv = ctx->priv_data;
-
-    av_freep(&priv->write_buffer);
-}
-
 const CodedBitstreamType ff_cbs_type_mpeg2 = {
     .codec_id          = AV_CODEC_ID_MPEG2VIDEO,
 
@@ -442,6 +390,4 @@ const CodedBitstreamType ff_cbs_type_mpeg2 = {
     .read_unit         = &cbs_mpeg2_read_unit,
     .write_unit        = &cbs_mpeg2_write_unit,
     .assemble_fragment = &cbs_mpeg2_assemble_fragment,
-
-    .close             = &cbs_mpeg2_close,
 };
diff --git a/libavcodec/cbs_mpeg2.h b/libavcodec/cbs_mpeg2.h
index 11f93b9df8..bdd3cb9301 100644
--- a/libavcodec/cbs_mpeg2.h
+++ b/libavcodec/cbs_mpeg2.h
@@ -219,10 +219,6 @@ typedef struct CodedBitstreamMPEG2Context {
     uint8_t scalable_mode;
     uint8_t progressive_sequence;
     uint8_t number_of_frame_centre_offsets;
-
-    // Write buffer.
-    uint8_t *write_buffer;
-    size_t write_buffer_size;
 } CodedBitstreamMPEG2Context;
 
 
diff --git a/libavcodec/cbs_vp9.c b/libavcodec/cbs_vp9.c
index 1c83257613..fa42d17982 100644
--- a/libavcodec/cbs_vp9.c
+++ b/libavcodec/cbs_vp9.c
@@ -528,62 +528,28 @@ static int cbs_vp9_read_unit(CodedBitstreamContext *ctx,
 }
 
 static int cbs_vp9_write_unit(CodedBitstreamContext *ctx,
-                              CodedBitstreamUnit *unit)
+                              CodedBitstreamUnit *unit,
+                              PutBitContext *pbc)
 {
-    CodedBitstreamVP9Context *priv = ctx->priv_data;
     VP9RawFrame *frame = unit->content;
-    PutBitContext pbc;
     int err;
 
-    if (!priv->write_buffer) {
-        // Initial write buffer size is 1MB.
-        priv->write_buffer_size = 1024 * 1024;
-
-    reallocate_and_try_again:
-        err = av_reallocp(&priv->write_buffer, priv->write_buffer_size);
-        if (err < 0) {
-            av_log(ctx->log_ctx, AV_LOG_ERROR, "Unable to allocate a "
-                   "sufficiently large write buffer (last attempt "
-                   "%"SIZE_SPECIFIER" bytes).\n", priv->write_buffer_size);
-            return err;
-        }
-    }
-
-    init_put_bits(&pbc, priv->write_buffer, priv->write_buffer_size);
-
-    err = cbs_vp9_write_frame(ctx, &pbc, frame);
-    if (err == AVERROR(ENOSPC)) {
-        priv->write_buffer_size *= 2;
-        goto reallocate_and_try_again;
-    }
+    err = cbs_vp9_write_frame(ctx, pbc, frame);
     if (err < 0)
         return err;
 
     // Frame must be byte-aligned.
-    av_assert0(put_bits_count(&pbc) % 8 == 0);
-
-    unit->data_size        = put_bits_count(&pbc) / 8;
-    unit->data_bit_padding = 0;
-    flush_put_bits(&pbc);
+    av_assert0(put_bits_count(pbc) % 8 == 0);
 
     if (frame->data) {
-        if (unit->data_size + frame->data_size >
-            priv->write_buffer_size) {
-            priv->write_buffer_size *= 2;
-            goto reallocate_and_try_again;
-        }
+        if (frame->data_size > put_bits_left(pbc) / 8)
+            return AVERROR(ENOSPC);
 
-        memcpy(priv->write_buffer + unit->data_size,
-               frame->data, frame->data_size);
-        unit->data_size += frame->data_size;
+        flush_put_bits(pbc);
+        memcpy(put_bits_ptr(pbc), frame->data, frame->data_size);
+        skip_put_bytes(pbc, frame->data_size);
     }
 
-    err = ff_cbs_alloc_unit_data(ctx, unit, unit->data_size);
-    if (err < 0)
-        return err;
-
-    memcpy(unit->data, priv->write_buffer, unit->data_size);
-
     return 0;
 }
 
@@ -677,13 +643,6 @@ static int cbs_vp9_assemble_fragment(CodedBitstreamContext *ctx,
     return 0;
 }
 
-static void cbs_vp9_close(CodedBitstreamContext *ctx)
-{
-    CodedBitstreamVP9Context *priv = ctx->priv_data;
-
-    av_freep(&priv->write_buffer);
-}
-
 const CodedBitstreamType ff_cbs_type_vp9 = {
     .codec_id          = AV_CODEC_ID_VP9,
 
@@ -693,6 +652,4 @@ const CodedBitstreamType ff_cbs_type_vp9 = {
     .read_unit         = &cbs_vp9_read_unit,
     .write_unit        = &cbs_vp9_write_unit,
     .assemble_fragment = &cbs_vp9_assemble_fragment,
-
-    .close             = &cbs_vp9_close,
 };
diff --git a/libavcodec/cbs_vp9.h b/libavcodec/cbs_vp9.h
index 4c9b2f880d..40e62476ed 100644
--- a/libavcodec/cbs_vp9.h
+++ b/libavcodec/cbs_vp9.h
@@ -207,10 +207,6 @@ typedef struct CodedBitstreamVP9Context {
     int bit_depth;
 
     VP9ReferenceFrameState ref[VP9_NUM_REF_FRAMES];
-
-    // Write buffer.
-    uint8_t *write_buffer;
-    size_t write_buffer_size;
 } CodedBitstreamVP9Context;