h2645_parse: Allocate a single buffer per packet

Message ID 20171103195245.9068-1-jamrial@gmail.com
State New
Headers show

Commit Message

James Almer Nov. 3, 2017, 7:52 p.m.
From: Kieran Kunhya <kierank@obe.tv>

Drastically reduces memory usage on pathological streams.
---
Fixed so it may apply cleanly, and with the memset() change removed.

 libavcodec/h2645_parse.c | 22 ++++++++++++----------
 libavcodec/h2645_parse.h | 10 ++++++++--
 libavcodec/h264_parser.c | 11 ++++++++---
 libavcodec/qsvenc_hevc.c | 13 +++++++++----
 4 files changed, 37 insertions(+), 19 deletions(-)

Comments

Luca Barbato Nov. 3, 2017, 10:04 p.m. | #1
On 03/11/2017 20:52, James Almer wrote:
> From: Kieran Kunhya <kierank@obe.tv>
> 
> Drastically reduces memory usage on pathological streams.
> ---
> Fixed so it may apply cleanly, and with the memset() change removed.

Thank you :)

I'm happy with it now, anybody is against landing it?

lu

Patch

diff --git a/libavcodec/h2645_parse.c b/libavcodec/h2645_parse.c
index b507b19ecb..2ee7672c05 100644
--- a/libavcodec/h2645_parse.c
+++ b/libavcodec/h2645_parse.c
@@ -30,7 +30,7 @@ 
 #include "h2645_parse.h"
 
 int ff_h2645_extract_rbsp(const uint8_t *src, int length,
-                          H2645NAL *nal)
+                          H2645RBSP *rbsp, H2645NAL *nal)
 {
     int i, si, di;
     uint8_t *dst;
@@ -88,11 +88,7 @@  int ff_h2645_extract_rbsp(const uint8_t *src, int length,
         return length;
     }
 
-    av_fast_malloc(&nal->rbsp_buffer, &nal->rbsp_buffer_size,
-                   length + AV_INPUT_BUFFER_PADDING_SIZE);
-    if (!nal->rbsp_buffer)
-        return AVERROR(ENOMEM);
-
+    nal->rbsp_buffer = &rbsp->rbsp_buffer[rbsp->rbsp_buffer_size];
     dst = nal->rbsp_buffer;
 
     memcpy(dst, src, i);
@@ -125,6 +121,8 @@  nsc:
     nal->size = di;
     nal->raw_data = src;
     nal->raw_size = si;
+    rbsp->rbsp_buffer_size += si;
+
     return si;
 }
 
@@ -220,6 +218,11 @@  int ff_h2645_packet_split(H2645Packet *pkt, const uint8_t *buf, int length,
     size_t next_avc = is_nalff ? 0 : length;
 
     bytestream2_init(&bc, buf, length);
+    av_fast_padded_malloc(&pkt->rbsp.rbsp_buffer, &pkt->rbsp.rbsp_buffer_alloc_size, length);
+    if (!pkt->rbsp.rbsp_buffer)
+        return AVERROR(ENOMEM);
+
+    pkt->rbsp.rbsp_buffer_size = 0;
 
     pkt->nb_nals = 0;
     while (bytestream2_get_bytes_left(&bc) >= 4) {
@@ -287,7 +290,7 @@  int ff_h2645_packet_split(H2645Packet *pkt, const uint8_t *buf, int length,
         }
         nal = &pkt->nals[pkt->nb_nals++];
 
-        consumed = ff_h2645_extract_rbsp(bc.buffer, extract_length, nal);
+        consumed = ff_h2645_extract_rbsp(bc.buffer, extract_length, &pkt->rbsp, nal);
         if (consumed < 0)
             return consumed;
 
@@ -322,9 +325,8 @@  int ff_h2645_packet_split(H2645Packet *pkt, const uint8_t *buf, int length,
 
 void ff_h2645_packet_uninit(H2645Packet *pkt)
 {
-    int i;
-    for (i = 0; i < pkt->nals_allocated; i++)
-        av_freep(&pkt->nals[i].rbsp_buffer);
     av_freep(&pkt->nals);
     pkt->nals_allocated = 0;
+    av_freep(&pkt->rbsp.rbsp_buffer);
+    pkt->rbsp.rbsp_buffer_alloc_size = pkt->rbsp.rbsp_buffer_size = 0;
 }
diff --git a/libavcodec/h2645_parse.h b/libavcodec/h2645_parse.h
index 9cc4441d9e..63ec493622 100644
--- a/libavcodec/h2645_parse.h
+++ b/libavcodec/h2645_parse.h
@@ -28,7 +28,6 @@ 
 
 typedef struct H2645NAL {
     uint8_t *rbsp_buffer;
-    int rbsp_buffer_size;
 
     int size;
     const uint8_t *data;
@@ -60,9 +59,16 @@  typedef struct H2645NAL {
     int ref_idc;
 } H2645NAL;
 
+typedef struct H2645RBSP {
+    uint8_t *rbsp_buffer;
+    int rbsp_buffer_alloc_size;
+    int rbsp_buffer_size;
+} H2645RBSP;
+
 /* an input packet split into unescaped NAL units */
 typedef struct H2645Packet {
     H2645NAL *nals;
+    H2645RBSP rbsp;
     int nb_nals;
     int nals_allocated;
 } H2645Packet;
@@ -71,7 +77,7 @@  typedef struct H2645Packet {
  * Extract the raw (unescaped) bitstream.
  */
 int ff_h2645_extract_rbsp(const uint8_t *src, int length,
-                          H2645NAL *nal);
+                          H2645RBSP *rbsp, H2645NAL *nal);
 
 /**
  * Split an input packet into NAL units.
diff --git a/libavcodec/h264_parser.c b/libavcodec/h264_parser.c
index 0bb78e09c8..710b4180f5 100644
--- a/libavcodec/h264_parser.c
+++ b/libavcodec/h264_parser.c
@@ -207,6 +207,7 @@  static inline int parse_nal_units(AVCodecParserContext *s,
     H264ParseContext *p = s->priv_data;
     const uint8_t *buf_end = buf + buf_size;
 
+    H2645RBSP rbsp = { NULL };
     H2645NAL nal = { NULL };
 
     unsigned int pps_id;
@@ -225,6 +226,10 @@  static inline int parse_nal_units(AVCodecParserContext *s,
     if (!buf_size)
         return 0;
 
+    av_fast_padded_malloc(&rbsp.rbsp_buffer, &rbsp.rbsp_buffer_alloc_size, buf_size);
+    if (!rbsp.rbsp_buffer)
+        return AVERROR(ENOMEM);
+
     for (;;) {
         const SPS *sps;
         int src_length, consumed;
@@ -250,7 +255,7 @@  static inline int parse_nal_units(AVCodecParserContext *s,
             break;
         }
 
-        consumed = ff_h2645_extract_rbsp(buf, src_length, &nal);
+        consumed = ff_h2645_extract_rbsp(buf, src_length, &rbsp, &nal);
         if (consumed < 0)
             break;
 
@@ -463,7 +468,7 @@  static inline int parse_nal_units(AVCodecParserContext *s,
                 s->field_order = AV_FIELD_UNKNOWN;
             }
 
-            av_freep(&nal.rbsp_buffer);
+            av_freep(&rbsp.rbsp_buffer);
             return 0; /* no need to evaluate the rest */
         }
         buf += consumed;
@@ -471,7 +476,7 @@  static inline int parse_nal_units(AVCodecParserContext *s,
     /* didn't find a picture! */
     av_log(avctx, AV_LOG_ERROR, "missing picture in access unit\n");
 fail:
-    av_freep(&nal.rbsp_buffer);
+    av_freep(&rbsp.rbsp_buffer);
     return -1;
 }
 
diff --git a/libavcodec/qsvenc_hevc.c b/libavcodec/qsvenc_hevc.c
index dbb55e2f18..0d9fa1b5dd 100644
--- a/libavcodec/qsvenc_hevc.c
+++ b/libavcodec/qsvenc_hevc.c
@@ -56,6 +56,7 @@  static int generate_fake_vps(QSVEncContext *q, AVCodecContext *avctx)
     PutByteContext pbc;
 
     GetBitContext gb;
+    H2645RBSP sps_rbsp = { NULL };
     H2645NAL sps_nal = { NULL };
     HEVCSPS sps = { 0 };
     HEVCVPS vps = { 0 };
@@ -69,8 +70,12 @@  static int generate_fake_vps(QSVEncContext *q, AVCodecContext *avctx)
         return AVERROR_UNKNOWN;
     }
 
+    av_fast_padded_malloc(&sps_rbsp.rbsp_buffer, &sps_rbsp.rbsp_buffer_alloc_size, avctx->extradata_size);
+    if (!sps_rbsp.rbsp_buffer)
+        return AVERROR(ENOMEM);
+
     /* parse the SPS */
-    ret = ff_h2645_extract_rbsp(avctx->extradata + 4, avctx->extradata_size - 4, &sps_nal);
+    ret = ff_h2645_extract_rbsp(avctx->extradata + 4, avctx->extradata_size - 4, &sps_rbsp, &sps_nal);
     if (ret < 0) {
         av_log(avctx, AV_LOG_ERROR, "Error unescaping the SPS buffer\n");
         return ret;
@@ -78,7 +83,7 @@  static int generate_fake_vps(QSVEncContext *q, AVCodecContext *avctx)
 
     ret = init_get_bits8(&gb, sps_nal.data, sps_nal.size);
     if (ret < 0) {
-        av_freep(&sps_nal.rbsp_buffer);
+        av_freep(&sps_rbsp.rbsp_buffer);
         return ret;
     }
 
@@ -87,13 +92,13 @@  static int generate_fake_vps(QSVEncContext *q, AVCodecContext *avctx)
     if (type != HEVC_NAL_SPS) {
         av_log(avctx, AV_LOG_ERROR, "Unexpected NAL type in the extradata: %d\n",
                type);
-        av_freep(&sps_nal.rbsp_buffer);
+        av_freep(&sps_rbsp.rbsp_buffer);
         return AVERROR_INVALIDDATA;
     }
     get_bits(&gb, 9);
 
     ret = ff_hevc_parse_sps(&sps, &gb, &sps_id, 0, NULL, avctx);
-    av_freep(&sps_nal.rbsp_buffer);
+    av_freep(&sps_rbsp.rbsp_buffer);
     if (ret < 0) {
         av_log(avctx, AV_LOG_ERROR, "Error parsing the SPS\n");
         return ret;