omxbufferpool: refactor to allow memory sharing

One big restriction of the OMX buffer pool has always been
that the GstMemory objects were flagged with NO_SHARE.
This was because the buffer pool needed to be sure that when
a buffer returned to the pool, it would be safe to release the
OMX buffer back to OpenMAX.

With this change, this is no longer a restriction. What this
commit introduces is a new allocator that allows us to track
the GstMemory objects independently. Now, when a buffer returns
to the pool, it is not necessary for the memory to be released
as well. We simply track the memory's ref count in the allocator
and we return the OMX buffer back when the memory's ref count
drops to 0.

The reason for doing this is to allow implementing zero-copy
transfers in situations where we may need to copy or map a
certain region of the buffer. For instance, omxh264enc ! h264parse
should be possible to be zero-copy by using an OMX buffer pool
between them.

7 files changed, 849 insertions, 265 deletions

diff --git a/omx/Makefile.am b/omx/Makefile.am
index 3019040..d8718f1 100644
--- a/omx/Makefile.am
+++ b/omx/Makefile.am
@@ -23,6 +23,7 @@ endif
 
 libgstomx_la_SOURCES = \
 	gstomx.c \
+	gstomxallocator.c \
 	gstomxbufferpool.c \
 	gstomxvideo.c \
 	gstomxvideodec.c \
@@ -53,6 +54,7 @@ libgstomx_la_SOURCES = \
 
 noinst_HEADERS = \
 	gstomx.h \
+	gstomxallocator.h \
 	gstomxbufferpool.h \
 	gstomxvideo.h \
 	gstomxvideodec.h \
diff --git a/omx/gstomxallocator.c b/omx/gstomxallocator.c
new file mode 100644
index 0000000..cf8058a
--- /dev/null
+++ b/omx/gstomxallocator.c
@@ -0,0 +1,552 @@
+/*
+ * Copyright (C) 2019, Collabora Ltd.
+ *   Author: George Kiagiadakis <george.kiagiadakis@collabora.com>
+ *
+ * This library is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation
+ * version 2.1 of the License.
+ *
+ * This library is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with this library; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301 USA
+ *
+ */
+
+#ifdef HAVE_CONFIG_H
+#include "config.h"
+#endif
+
+#include "gstomxallocator.h"
+#include <gst/allocators/gstdmabuf.h>
+
+GST_DEBUG_CATEGORY_STATIC (gst_omx_allocator_debug_category);
+#define GST_CAT_DEFAULT gst_omx_allocator_debug_category
+
+#define DEBUG_INIT \
+  GST_DEBUG_CATEGORY_INIT (gst_omx_allocator_debug_category, "omxallocator", 0, \
+      "debug category for gst-omx allocator class");
+
+G_DEFINE_TYPE_WITH_CODE (GstOMXAllocator, gst_omx_allocator, GST_TYPE_ALLOCATOR,
+    DEBUG_INIT);
+
+enum
+{
+  SIG_OMXBUF_RELEASED,
+  SIG_FOREIGN_MEM_RELEASED,
+  LAST_SIGNAL
+};
+
+static guint signals[LAST_SIGNAL] = { 0 };
+
+/* Custom allocator for memory associated with OpenMAX buffers
+ *
+ * The main purpose of this allocator is to track memory that is associated
+ * with OpenMAX buffers, so that we know when the buffers can be released
+ * back to OpenMAX.
+ *
+ * This allocator looks and behaves more like a buffer pool. It allocates
+ * the memory objects before starting and sets a miniobject dispose function
+ * on them, which allows them to return when their last ref count is dropped.
+ *
+ * The type of memory that this allocator manages is GstOMXMemory. However, it
+ * is possible to manage a different type of memory, in which case the
+ * GstOMXMemory object is used only internally. There are two supported cases:
+ * - Allocate memory from the dmabuf allocator
+ * - Take memory that was allocated externally and manage it here
+ *
+ * In both cases, this allocator will replace the miniobject dispose function
+ * of these memory objects, so if they were acquired from here, they will also
+ * return here on their last unref.
+ *
+ * The caller initially needs to configure how many memory objects will be
+ * managed here by calling configure(). After that it needs to call
+ * set_active(TRUE) and finally allocate() for each memory. Allocation is done
+ * like this to facilitate calling allocate() from the alloc() function of
+ * the buffer pool for each OMX buffer on the port.
+ *
+ * After the allocator has been activated and all buffers have been allocated,
+ * the acquire() method can be called to retrieve a memory object. acquire() can
+ * be given an OMX buffer index or pointer to locate and return the memory
+ * object that corresponds to this OMX buffer. If the buffer is already
+ * acquired, this will result in a GST_FLOW_ERROR.
+ *
+ * When the last reference count is dropped on a memory that was acquired from
+ * here, its dispose function will ref it again and allow it to be acquired
+ * again. In addition, the omxbuf-released signal is fired to let the caller
+ * know that it can return this OMX buffer to the port, as it is no longer
+ * used outside this allocator.
+ */
+
+/******************/
+/** GstOMXMemory **/
+/******************/
+
+#define GST_OMX_MEMORY_TYPE "openmax"
+
+GQuark
+gst_omx_memory_quark (void)
+{
+  static GQuark quark = 0;
+
+  if (quark == 0)
+    quark = g_quark_from_static_string ("GstOMXMemory");
+
+  return quark;
+}
+
+static GstOMXMemory *
+gst_omx_memory_new (GstOMXAllocator * allocator, GstOMXBuffer * omx_buf,
+    GstMemoryFlags flags, GstMemory * parent, gssize offset, gssize size)
+{
+  GstOMXMemory *mem;
+  gint align;
+  gsize maxsize;
+
+  /* GStreamer uses a bitmask for the alignment while
+   * OMX uses the alignment itself. So we have to convert
+   * here */
+  align = allocator->port->port_def.nBufferAlignment;
+  if (align > 0)
+    align -= 1;
+  if (((align + 1) & align) != 0) {
+    GST_WARNING ("Invalid alignment that is not a power of two: %u",
+        (guint) allocator->port->port_def.nBufferAlignment);
+    align = 0;
+  }
+
+  maxsize = omx_buf->omx_buf->nAllocLen;
+
+  if (size == -1) {
+    size = maxsize - offset;
+  }
+
+  mem = g_slice_new0 (GstOMXMemory);
+  gst_memory_init (GST_MEMORY_CAST (mem), flags, (GstAllocator *) allocator,
+      parent, maxsize, align, offset, size);
+
+  mem->buf = omx_buf;
+
+  return mem;
+}
+
+static gpointer
+gst_omx_memory_map (GstMemory * mem, gsize maxsize, GstMapFlags flags)
+{
+  GstOMXMemory *omem = (GstOMXMemory *) mem;
+
+  /* if we are using foreign_mem, the GstOMXMemory should never appear
+   * anywhere outside this allocator, therefore it should never be mapped */
+  g_return_val_if_fail (!omem->foreign_mem, NULL);
+
+  return omem->buf->omx_buf->pBuffer;
+}
+
+static void
+gst_omx_memory_unmap (GstMemory * mem)
+{
+}
+
+static GstMemory *
+gst_omx_memory_share (GstMemory * mem, gssize offset, gssize size)
+{
+  GstOMXMemory *omem = (GstOMXMemory *) mem;
+  GstOMXMemory *sub;
+  GstMemory *parent;
+
+  /* find the real parent */
+  if ((parent = mem->parent) == NULL)
+    parent = mem;
+
+  if (size == -1)
+    size = mem->size - offset;
+
+  /* the shared memory is always readonly */
+  sub = gst_omx_memory_new ((GstOMXAllocator *) mem->allocator, omem->buf,
+      GST_MINI_OBJECT_FLAGS (parent) | GST_MINI_OBJECT_FLAG_LOCK_READONLY,
+      parent, offset, size);
+
+  return (GstMemory *) sub;
+}
+
+GstOMXBuffer *
+gst_omx_memory_get_omx_buf (GstMemory * mem)
+{
+  GstOMXMemory *omx_mem;
+
+  if (GST_IS_OMX_ALLOCATOR (mem->allocator))
+    omx_mem = (GstOMXMemory *) mem;
+  else
+    omx_mem = gst_mini_object_get_qdata (GST_MINI_OBJECT (mem),
+        GST_OMX_MEMORY_QUARK);
+
+  if (!omx_mem)
+    return NULL;
+
+  return omx_mem->buf;
+}
+
+/*********************/
+/** GstOMXAllocator **/
+/*********************/
+
+static void
+gst_omx_allocator_init (GstOMXAllocator * allocator)
+{
+  GstAllocator *alloc = GST_ALLOCATOR_CAST (allocator);
+
+  alloc->mem_type = GST_OMX_MEMORY_TYPE;
+
+  alloc->mem_map = gst_omx_memory_map;
+  alloc->mem_unmap = gst_omx_memory_unmap;
+  alloc->mem_share = gst_omx_memory_share;
+  /* default copy & is_span */
+
+  GST_OBJECT_FLAG_SET (allocator, GST_ALLOCATOR_FLAG_CUSTOM_ALLOC);
+
+  g_mutex_init (&allocator->lock);
+  g_cond_init (&allocator->cond);
+}
+
+GstOMXAllocator *
+gst_omx_allocator_new (GstOMXComponent * component, GstOMXPort * port)
+{
+  GstOMXAllocator *allocator;
+
+  allocator = g_object_new (gst_omx_allocator_get_type (), NULL);
+  allocator->component = gst_omx_component_ref (component);
+  allocator->port = port;
+
+  return allocator;
+}
+
+static void
+gst_omx_allocator_finalize (GObject * object)
+{
+  GstOMXAllocator *allocator = GST_OMX_ALLOCATOR (object);
+
+  gst_omx_component_unref (allocator->component);
+  g_mutex_clear (&allocator->lock);
+  g_cond_clear (&allocator->cond);
+}
+
+gboolean
+gst_omx_allocator_configure (GstOMXAllocator * allocator, guint count,
+    GstOMXAllocatorForeignMemMode mode)
+{
+  /* check if already configured */
+  if (allocator->n_memories > 0)
+    return FALSE;
+
+  allocator->n_memories = count;
+  allocator->foreign_mode = mode;
+  if (mode == GST_OMX_ALLOCATOR_FOREIGN_MEM_DMABUF)
+    allocator->foreign_allocator = gst_dmabuf_allocator_new ();
+
+  return TRUE;
+}
+
+/* must be protected with allocator->lock */
+static void
+gst_omx_allocator_dealloc (GstOMXAllocator * allocator)
+{
+  /* might be called more than once */
+  if (!allocator->memories)
+    return;
+
+  /* return foreign memory back to whoever lended it to us.
+   * the signal handler is expected to increase the ref count of foreign_mem */
+  if (allocator->foreign_mode == GST_OMX_ALLOCATOR_FOREIGN_MEM_OTHER_POOL) {
+    gint i;
+    GstOMXMemory *m;
+
+    for (i = 0; i < allocator->memories->len; i++) {
+      m = g_ptr_array_index (allocator->memories, i);
+
+      /* this should not happen, but let's not crash for this */
+      if (!m->foreign_mem) {
+        GST_WARNING_OBJECT (allocator, "no foreign_mem to release");
+        continue;
+      }
+
+      /* restore the original dispose function */
+      GST_MINI_OBJECT_CAST (m->foreign_mem)->dispose =
+          (GstMiniObjectDisposeFunction) m->foreign_dispose;
+
+      g_signal_emit (allocator, signals[SIG_FOREIGN_MEM_RELEASED], 0, i,
+          m->foreign_mem);
+    }
+  }
+
+  g_ptr_array_foreach (allocator->memories, (GFunc) gst_memory_unref, NULL);
+  g_ptr_array_free (allocator->memories, TRUE);
+  allocator->memories = NULL;
+  allocator->n_memories = 0;
+  allocator->foreign_mode = GST_OMX_ALLOCATOR_FOREIGN_MEM_NONE;
+  if (allocator->foreign_allocator) {
+    g_object_unref (allocator->foreign_allocator);
+    allocator->foreign_allocator = NULL;
+  }
+
+  g_cond_broadcast (&allocator->cond);
+}
+
+gboolean
+gst_omx_allocator_set_active (GstOMXAllocator * allocator, gboolean active)
+{
+  gboolean changed = FALSE;
+
+  /* on activation, _configure() must be called first */
+  g_return_val_if_fail (!active || allocator->n_memories > 0, FALSE);
+
+  g_mutex_lock (&allocator->lock);
+
+  if (allocator->active != active)
+    changed = TRUE;
+
+  if (changed) {
+    if (active) {
+      allocator->memories = g_ptr_array_sized_new (allocator->n_memories);
+      g_ptr_array_set_size (allocator->memories, allocator->n_memories);
+    } else {
+      if (g_atomic_int_get (&allocator->n_outstanding) == 0)
+        gst_omx_allocator_dealloc (allocator);
+    }
+  }
+
+  allocator->active = active;
+  g_mutex_unlock (&allocator->lock);
+
+  return changed;
+}
+
+void
+gst_omx_allocator_wait_inactive (GstOMXAllocator * allocator)
+{
+  g_mutex_lock (&allocator->lock);
+  while (allocator->memories)
+    g_cond_wait (&allocator->cond, &allocator->lock);
+  g_mutex_unlock (&allocator->lock);
+}
+
+static inline void
+dec_outstanding (GstOMXAllocator * allocator)
+{
+  if (g_atomic_int_dec_and_test (&allocator->n_outstanding)) {
+    /* keep a ref to the allocator because _dealloc() will free
+     * all the memories and the memories might be the only thing holding
+     * a reference to the allocator; we need to keep it alive until the
+     * end of this function call */
+    g_object_ref (allocator);
+
+    /* take the lock so that _set_active() is not run concurrently */
+    g_mutex_lock (&allocator->lock);
+
+    /* now that we have the lock, check if we have been de-activated with
+     * outstanding buffers */
+    if (!allocator->active)
+      gst_omx_allocator_dealloc (allocator);
+
+    g_mutex_unlock (&allocator->lock);
+    g_object_unref (allocator);
+  }
+}
+
+GstFlowReturn
+gst_omx_allocator_acquire (GstOMXAllocator * allocator, GstMemory ** memory,
+    gint index, GstOMXBuffer * omx_buf)
+{
+  GstFlowReturn ret = GST_FLOW_OK;
+  GstOMXMemory *omx_mem = NULL;
+
+  /* ensure memories are not going to disappear concurrently */
+  g_atomic_int_inc (&allocator->n_outstanding);
+
+  if (!allocator->active) {
+    ret = GST_FLOW_FLUSHING;
+    goto beach;
+  }
+
+  if (index >= 0 && index < allocator->n_memories)
+    omx_mem = g_ptr_array_index (allocator->memories, index);
+  else if (omx_buf) {
+    for (index = 0; index < allocator->n_memories; index++) {
+      omx_mem = g_ptr_array_index (allocator->memories, index);
+      if (omx_mem->buf == omx_buf)
+        break;
+    }
+  }
+
+  if (G_UNLIKELY (!omx_mem || index >= allocator->n_memories)) {
+    GST_ERROR_OBJECT (allocator, "Failed to find OMX memory");
+    ret = GST_FLOW_ERROR;
+    goto beach;
+  }
+
+  if (G_UNLIKELY (omx_mem->buf->used)) {
+    GST_ERROR_OBJECT (allocator,
+        "Trying to acquire a buffer that is being used by the OMX port");
+    ret = GST_FLOW_ERROR;
+    goto beach;
+  }
+
+  omx_mem->acquired = TRUE;
+
+  if (omx_mem->foreign_mem)
+    *memory = omx_mem->foreign_mem;
+  else
+    *memory = GST_MEMORY_CAST (omx_mem);
+
+beach:
+  if (ret != GST_FLOW_OK)
+    dec_outstanding (allocator);
+  return ret;
+}
+
+/* installed as the GstMiniObject::dispose function of the acquired GstMemory */
+static gboolean
+gst_omx_allocator_memory_dispose (GstMemory * mem)
+{
+  GstOMXMemory *omx_mem;
+  GstOMXAllocator *allocator;
+
+  /* memory may be from our allocator, but
+   * may as well be from the dmabuf allocator */
+  if (GST_IS_OMX_ALLOCATOR (mem->allocator))
+    omx_mem = (GstOMXMemory *) mem;
+  else
+    omx_mem = gst_mini_object_get_qdata (GST_MINI_OBJECT (mem),
+        GST_OMX_MEMORY_QUARK);
+
+  if (omx_mem->acquired) {
+    /* keep the memory alive */
+    gst_memory_ref (mem);
+
+    omx_mem->acquired = FALSE;
+
+    allocator = GST_OMX_ALLOCATOR (GST_MEMORY_CAST (omx_mem)->allocator);
+
+    /* inform the upper layer that we are no longer using this GstOMXBuffer */
+    g_signal_emit (allocator, signals[SIG_OMXBUF_RELEASED], 0, omx_mem->buf);
+
+    dec_outstanding (allocator);
+
+    /* be careful here, both the memory and the allocator
+     * may have been free'd as part of the call to dec_outstanding() */
+
+    return FALSE;
+  }
+
+  /* if the foreign memory had a dispose function, let that one decide
+   * the fate of this memory. We are no longer going to be using it here */
+  if (omx_mem->foreign_dispose)
+    return omx_mem->foreign_dispose (GST_MINI_OBJECT_CAST (mem));
+
+  return TRUE;
+}
+
+static inline void
+install_mem_dispose (GstOMXMemory * mem)
+{
+  GstMemory *managed_mem = (GstMemory *) mem;
+
+  if (mem->foreign_mem) {
+    managed_mem = mem->foreign_mem;
+    mem->foreign_dispose = GST_MINI_OBJECT_CAST (managed_mem)->dispose;
+  }
+
+  GST_MINI_OBJECT_CAST (managed_mem)->dispose =
+      (GstMiniObjectDisposeFunction) gst_omx_allocator_memory_dispose;
+}
+
+/* the returned memory is transfer:none, ref still belongs to the allocator */
+GstMemory *
+gst_omx_allocator_allocate (GstOMXAllocator * allocator, gint index,
+    GstMemory * foreign_mem)
+{
+  GstOMXMemory *mem;
+  GstOMXBuffer *omx_buf;
+
+  g_return_val_if_fail (allocator->port->buffers, NULL);
+  g_return_val_if_fail (allocator->memories, NULL);
+  g_return_val_if_fail (index >= 0 && index < allocator->n_memories, NULL);
+  g_return_val_if_fail ((foreign_mem == NULL &&
+          allocator->foreign_mode != GST_OMX_ALLOCATOR_FOREIGN_MEM_OTHER_POOL)
+      || (foreign_mem != NULL
+          && allocator->foreign_mode ==
+          GST_OMX_ALLOCATOR_FOREIGN_MEM_OTHER_POOL), NULL);
+
+  omx_buf = g_ptr_array_index (allocator->port->buffers, index);
+  g_return_val_if_fail (omx_buf != NULL, NULL);
+
+  mem = gst_omx_memory_new (allocator, omx_buf, 0, NULL, 0, -1);
+
+  switch (allocator->foreign_mode) {
+    case GST_OMX_ALLOCATOR_FOREIGN_MEM_NONE:
+      install_mem_dispose (mem);
+      break;
+    case GST_OMX_ALLOCATOR_FOREIGN_MEM_DMABUF:
+    {
+      gint fd = GPOINTER_TO_INT (omx_buf->omx_buf->pBuffer);
+      mem->foreign_mem =
+          gst_dmabuf_allocator_alloc (allocator->foreign_allocator, fd,
+          omx_buf->omx_buf->nAllocLen);
+      gst_mini_object_set_qdata (GST_MINI_OBJECT (mem->foreign_mem),
+          GST_OMX_MEMORY_QUARK, mem, NULL);
+      install_mem_dispose (mem);
+      break;
+    }
+    case GST_OMX_ALLOCATOR_FOREIGN_MEM_OTHER_POOL:
+      mem->foreign_mem = foreign_mem;
+      gst_mini_object_set_qdata (GST_MINI_OBJECT (mem->foreign_mem),
+          GST_OMX_MEMORY_QUARK, mem, NULL);
+      install_mem_dispose (mem);
+      break;
+    default:
+      g_assert_not_reached ();
+      break;
+  }
+
+  g_ptr_array_index (allocator->memories, index) = mem;
+  return mem->foreign_mem ? mem->foreign_mem : (GstMemory *) mem;
+}
+
+static void
+gst_omx_allocator_free (GstAllocator * allocator, GstMemory * mem)
+{
+  GstOMXMemory *omem = (GstOMXMemory *) mem;
+
+  g_warn_if_fail (!omem->acquired);
+
+  if (omem->foreign_mem)
+    gst_memory_unref (omem->foreign_mem);
+
+  g_slice_free (GstOMXMemory, omem);
+}
+
+static void
+gst_omx_allocator_class_init (GstOMXAllocatorClass * klass)
+{
+  GObjectClass *object_class;
+  GstAllocatorClass *allocator_class;
+
+  object_class = (GObjectClass *) klass;
+  allocator_class = (GstAllocatorClass *) klass;
+
+  object_class->finalize = gst_omx_allocator_finalize;
+  allocator_class->alloc = NULL;
+  allocator_class->free = gst_omx_allocator_free;
+
+  signals[SIG_OMXBUF_RELEASED] = g_signal_new ("omxbuf-released",
+      G_TYPE_FROM_CLASS (klass), G_SIGNAL_RUN_LAST, 0,
+      NULL, NULL, NULL, G_TYPE_NONE, 1, G_TYPE_POINTER);
+
+  signals[SIG_FOREIGN_MEM_RELEASED] = g_signal_new ("foreign-mem-released",
+      G_TYPE_FROM_CLASS (klass), G_SIGNAL_RUN_LAST, 0,
+      NULL, NULL, NULL, G_TYPE_NONE, 2, G_TYPE_INT, G_TYPE_POINTER);
+}
diff --git a/omx/gstomxallocator.h b/omx/gstomxallocator.h
new file mode 100644
index 0000000..5c94584
--- /dev/null
+++ b/omx/gstomxallocator.h
@@ -0,0 +1,114 @@
+/*
+ * Copyright (C) 2019, Collabora Ltd.
+ *   Author: George Kiagiadakis <george.kiagiadakis@collabora.com>
+ *
+ * This library is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation
+ * version 2.1 of the License.
+ *
+ * This library is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with this library; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301 USA
+ *
+ */
+
+#ifndef __GST_OMX_ALLOCATOR_H__
+#define __GST_OMX_ALLOCATOR_H__
+
+#ifdef HAVE_CONFIG_H
+#include "config.h"
+#endif
+
+#include <gst/gst.h>
+
+#include "gstomx.h"
+
+G_BEGIN_DECLS
+
+#define GST_OMX_MEMORY_QUARK gst_omx_memory_quark ()
+
+#define GST_TYPE_OMX_ALLOCATOR   (gst_omx_allocator_get_type())
+#define GST_IS_OMX_ALLOCATOR(obj) (G_TYPE_CHECK_INSTANCE_TYPE ((obj), GST_TYPE_OMX_ALLOCATOR))
+#define GST_OMX_ALLOCATOR(obj)    (G_TYPE_CHECK_INSTANCE_CAST ((obj), GST_TYPE_OMX_ALLOCATOR, GstOMXAllocator))
+
+typedef struct _GstOMXMemory GstOMXMemory;
+typedef struct _GstOMXAllocator GstOMXAllocator;
+typedef struct _GstOMXAllocatorClass GstOMXAllocatorClass;
+
+typedef enum {
+  GST_OMX_ALLOCATOR_FOREIGN_MEM_NONE,
+  GST_OMX_ALLOCATOR_FOREIGN_MEM_DMABUF,
+  GST_OMX_ALLOCATOR_FOREIGN_MEM_OTHER_POOL,
+} GstOMXAllocatorForeignMemMode;
+
+struct _GstOMXMemory
+{
+  GstMemory mem;
+  GstOMXBuffer *buf;
+
+  /* TRUE if the memory is in use outside the allocator */
+  gboolean acquired;
+
+  /* memory allocated from the foreign_allocator
+   * or planted externally when using a foreign buffer pool */
+  GstMemory *foreign_mem;
+  /* the original dispose function of foreign_mem */
+  GstMiniObjectDisposeFunction foreign_dispose;
+};
+
+struct _GstOMXAllocator
+{
+  GstAllocator parent;
+
+  GstOMXComponent *component;
+  GstOMXPort *port;
+
+  GstOMXAllocatorForeignMemMode foreign_mode;
+  GstAllocator *foreign_allocator;
+
+  /* array of GstOMXMemory */
+  GPtrArray *memories;
+  guint n_memories;
+
+  guint n_outstanding;
+  gboolean active;
+
+  GMutex lock;
+  GCond cond;
+};
+
+struct _GstOMXAllocatorClass
+{
+  GstAllocatorClass parent_class;
+};
+
+GType gst_omx_allocator_get_type (void);
+
+GQuark gst_omx_memory_quark (void);
+
+GstOMXBuffer * gst_omx_memory_get_omx_buf (GstMemory * mem);
+
+GstOMXAllocator * gst_omx_allocator_new (GstOMXComponent * component,
+    GstOMXPort * port);
+
+gboolean gst_omx_allocator_configure (GstOMXAllocator * allocator, guint count,
+    GstOMXAllocatorForeignMemMode mode);
+gboolean gst_omx_allocator_set_active (GstOMXAllocator * allocator,
+    gboolean active);
+void gst_omx_allocator_wait_inactive (GstOMXAllocator * allocator);
+
+GstFlowReturn gst_omx_allocator_acquire (GstOMXAllocator * allocator,
+    GstMemory ** memory, gint index, GstOMXBuffer * omx_buf);
+
+GstMemory * gst_omx_allocator_allocate (GstOMXAllocator * allocator, gint index,
+    GstMemory * foreign_mem);
+
+G_END_DECLS
+
+#endif
diff --git a/omx/gstomxbufferpool.c b/omx/gstomxbufferpool.c
index a77ade6..094f4c5 100644
--- a/omx/gstomxbufferpool.c
+++ b/omx/gstomxbufferpool.c
@@ -1,8 +1,9 @@
 /*
  * Copyright (C) 2011, Hewlett-Packard Development Company, L.P.
  *   Author: Sebastian Dröge <sebastian.droege@collabora.co.uk>, Collabora Ltd.
- * Copyright (C) 2013, Collabora Ltd.
+ * Copyright (C) 2013-2019, Collabora Ltd.
  *   Author: Sebastian Dröge <sebastian.droege@collabora.co.uk>
+ *           George Kiagiadakis <george.kiagiadakis@collabora.com>
  *
  * This library is free software; you can redistribute it and/or
  * modify it under the terms of the GNU Lesser General Public
@@ -31,10 +32,6 @@
 GST_DEBUG_CATEGORY_STATIC (gst_omx_buffer_pool_debug_category);
 #define GST_CAT_DEFAULT gst_omx_buffer_pool_debug_category
 
-typedef struct _GstOMXMemory GstOMXMemory;
-typedef struct _GstOMXMemoryAllocator GstOMXMemoryAllocator;
-typedef struct _GstOMXMemoryAllocatorClass GstOMXMemoryAllocatorClass;
-
 enum
 {
   SIG_ALLOCATE,
@@ -43,135 +40,6 @@ enum
 
 static guint signals[LAST_SIGNAL] = { 0 };
 
-struct _GstOMXMemory
-{
-  GstMemory mem;
-
-  GstOMXBuffer *buf;
-};
-
-struct _GstOMXMemoryAllocator
-{
-  GstAllocator parent;
-};
-
-struct _GstOMXMemoryAllocatorClass
-{
-  GstAllocatorClass parent_class;
-};
-
-#define GST_OMX_MEMORY_TYPE "openmax"
-
-static GstMemory *
-gst_omx_memory_allocator_alloc_dummy (GstAllocator * allocator, gsize size,
-    GstAllocationParams * params)
-{
-  g_assert_not_reached ();
-  return NULL;
-}
-
-static void
-gst_omx_memory_allocator_free (GstAllocator * allocator, GstMemory * mem)
-{
-  GstOMXMemory *omem = (GstOMXMemory *) mem;
-
-  /* TODO: We need to remember which memories are still used
-   * so we can wait until everything is released before allocating
-   * new memory
-   */
-
-  g_slice_free (GstOMXMemory, omem);
-}
-
-static gpointer
-gst_omx_memory_map (GstMemory * mem, gsize maxsize, GstMapFlags flags)
-{
-  GstOMXMemory *omem = (GstOMXMemory *) mem;
-
-  return omem->buf->omx_buf->pBuffer;
-}
-
-static void
-gst_omx_memory_unmap (GstMemory * mem)
-{
-}
-
-static GstMemory *
-gst_omx_memory_share (GstMemory * mem, gssize offset, gssize size)
-{
-  g_assert_not_reached ();
-  return NULL;
-}
-
-GType gst_omx_memory_allocator_get_type (void);
-G_DEFINE_TYPE (GstOMXMemoryAllocator, gst_omx_memory_allocator,
-    GST_TYPE_ALLOCATOR);
-
-#define GST_TYPE_OMX_MEMORY_ALLOCATOR   (gst_omx_memory_allocator_get_type())
-#define GST_IS_OMX_MEMORY_ALLOCATOR(obj) (G_TYPE_CHECK_INSTANCE_TYPE ((obj), GST_TYPE_OMX_MEMORY_ALLOCATOR))
-
-static void
-gst_omx_memory_allocator_class_init (GstOMXMemoryAllocatorClass * klass)
-{
-  GstAllocatorClass *allocator_class;
-
-  allocator_class = (GstAllocatorClass *) klass;
-
-  allocator_class->alloc = gst_omx_memory_allocator_alloc_dummy;
-  allocator_class->free = gst_omx_memory_allocator_free;
-}
-
-static void
-gst_omx_memory_allocator_init (GstOMXMemoryAllocator * allocator)
-{
-  GstAllocator *alloc = GST_ALLOCATOR_CAST (allocator);
-
-  alloc->mem_type = GST_OMX_MEMORY_TYPE;
-  alloc->mem_map = gst_omx_memory_map;
-  alloc->mem_unmap = gst_omx_memory_unmap;
-  alloc->mem_share = gst_omx_memory_share;
-
-  /* default copy & is_span */
-
-  GST_OBJECT_FLAG_SET (allocator, GST_ALLOCATOR_FLAG_CUSTOM_ALLOC);
-}
-
-static GstMemory *
-gst_omx_memory_allocator_alloc (GstAllocator * allocator, GstMemoryFlags flags,
-    GstOMXBuffer * buf)
-{
-  GstOMXMemory *mem;
-  gint align;
-
-  /* FIXME: We don't allow sharing because we need to know
-   * when the memory becomes unused and can only then put
-   * it back to the pool. Which is done in the pool's release
-   * function
-   */
-  flags |= GST_MEMORY_FLAG_NO_SHARE;
-
-  /* GStreamer uses a bitmask for the alignment while
-   * OMX uses the alignment itself. So we have to convert
-   * here */
-  align = buf->port->port_def.nBufferAlignment;
-  if (align > 0)
-    align -= 1;
-  if (((align + 1) & align) != 0) {
-    GST_WARNING ("Invalid alignment that is not a power of two: %u",
-        (guint) buf->port->port_def.nBufferAlignment);
-    align = 0;
-  }
-
-  mem = g_slice_new (GstOMXMemory);
-  /* the shared memory is always readonly */
-  gst_memory_init (GST_MEMORY_CAST (mem), flags, allocator, NULL,
-      buf->omx_buf->nAllocLen, align, 0, buf->omx_buf->nAllocLen);
-
-  mem->buf = buf;
-
-  return GST_MEMORY_CAST (mem);
-}
-
 /* Buffer pool for the buffers of an OpenMAX port.
  *
  * This pool is only used if we either passed buffers from another
@@ -211,6 +79,23 @@ gst_omx_memory_allocator_alloc (GstAllocator * allocator, GstMemoryFlags flags,
  *
  * For buffers provided to downstream, the buffer will be returned
  * back to the component (OMX_FillThisBuffer()) when it is released.
+ *
+ * This pool uses a special allocator object, GstOMXAllocator. The main purpose
+ * of this allocator is to track GstMemory objects in the same way that a
+ * GstBufferPool tracks buffers. When a buffer is inserted into this pool
+ * (either because it was just allocated or because it was released back to
+ * the pool), its memory is ripped off and is tracked separately by the
+ * allocator. When a buffer is then acquired, we acquire the corresponding
+ * GstMemory from the allocator and put it back in the buffer.
+ *
+ * This allocator mechanism allows us to track memory that has been shared
+ * with buffers that are not part of this pool. When a memory is shared, then
+ * its ref count is > 1, which means it will not be released to the allocator
+ * until the sub-memory is destroyed.
+ *
+ * When a memory returns to the allocator, the allocator fires the
+ * omxbuf-released signal, which is handled by the buffer pool to return the
+ * omx buffer to the port or the queue.
  */
 
 #define DEBUG_INIT \
@@ -230,6 +115,7 @@ gst_omx_buffer_pool_start (GstBufferPool * bpool)
   gboolean has_buffers;
   GstStructure *config;
   guint min, max;
+  GstOMXAllocatorForeignMemMode mode;
 
   /* Only allow to start the pool if we still are attached
    * to a component and port */
@@ -271,6 +157,19 @@ gst_omx_buffer_pool_start (GstBufferPool * bpool)
 
   g_assert (pool->port->buffers);
 
+  if (pool->output_mode == GST_OMX_BUFFER_MODE_DMABUF)
+    mode = GST_OMX_ALLOCATOR_FOREIGN_MEM_DMABUF;
+  else if (pool->other_pool)
+    mode = GST_OMX_ALLOCATOR_FOREIGN_MEM_OTHER_POOL;
+  else
+    mode = GST_OMX_ALLOCATOR_FOREIGN_MEM_NONE;
+
+  if (!gst_omx_allocator_configure (pool->allocator, min, mode))
+    return FALSE;
+
+  if (!gst_omx_allocator_set_active (pool->allocator, TRUE))
+    return FALSE;
+
   return
       GST_BUFFER_POOL_CLASS (gst_omx_buffer_pool_parent_class)->start (bpool);
 }
@@ -279,23 +178,18 @@ static gboolean
 gst_omx_buffer_pool_stop (GstBufferPool * bpool)
 {
   GstOMXBufferPool *pool = GST_OMX_BUFFER_POOL (bpool);
-  gint i = 0;
-
-  if (pool->buffers) {
-    /* When not using the default GstBufferPool::GstAtomicQueue then
-     * GstBufferPool::free_buffer is not called while stopping the pool
-     * (because the queue is empty) */
-    for (i = 0; i < pool->buffers->len; i++)
-      GST_BUFFER_POOL_CLASS (gst_omx_buffer_pool_parent_class)->release_buffer
-          (bpool, g_ptr_array_index (pool->buffers, i));
-
-    /* Remove any buffers that are there */
-    g_ptr_array_set_size (pool->buffers, 0);
-  }
 
   /* Remove any buffers that are there */
   g_ptr_array_set_size (pool->buffers, 0);
 
+  GST_DEBUG_OBJECT (pool, "deactivating OMX allocator");
+  gst_omx_allocator_set_active (pool->allocator, FALSE);
+
+  /* ensure all memories have been deallocated;
+   * this may take a while if some memories are being shared
+   * and therefore are in use somewhere else in the pipeline */
+  gst_omx_allocator_wait_inactive (pool->allocator);
+
   GST_DEBUG_OBJECT (pool, "deallocate OMX buffers");
   gst_omx_port_deallocate_buffers (pool->port);
 
@@ -336,6 +230,8 @@ gst_omx_buffer_pool_set_config (GstBufferPool * bpool, GstStructure * config)
   GstOMXBufferPool *pool = GST_OMX_BUFFER_POOL (bpool);
   GstCaps *caps;
   guint size, min;
+  GstStructure *fake_config;
+  gboolean ret;
 
   GST_OBJECT_LOCK (pool);
 
@@ -371,8 +267,17 @@ gst_omx_buffer_pool_set_config (GstBufferPool * bpool, GstStructure * config)
 
   GST_OBJECT_UNLOCK (pool);
 
-  return GST_BUFFER_POOL_CLASS (gst_omx_buffer_pool_parent_class)->set_config
-      (bpool, config);
+  /* give a fake config to the parent default_set_config() with size == 0
+   * this prevents default_release_buffer() from free'ing the buffers, since
+   * we release them with no memory */
+  fake_config = gst_structure_copy (config);
+  gst_buffer_pool_config_set_params (fake_config, caps, 0, min, min);
+
+  ret = GST_BUFFER_POOL_CLASS (gst_omx_buffer_pool_parent_class)->set_config
+      (bpool, fake_config);
+  gst_structure_free (fake_config);
+
+  return ret;
 
   /* ERRORS */
 wrong_config:
@@ -402,29 +307,23 @@ gst_omx_buffer_pool_alloc_buffer (GstBufferPool * bpool,
 {
   GstOMXBufferPool *pool = GST_OMX_BUFFER_POOL (bpool);
   GstBuffer *buf;
-  GstOMXBuffer *omx_buf;
-
-  omx_buf = g_ptr_array_index (pool->port->buffers, pool->current_buffer_index);
-  g_return_val_if_fail (omx_buf != NULL, GST_FLOW_ERROR);
+  GstMemory *mem;
+  GstMemory *foreign_mem = NULL;
 
   if (pool->other_pool) {
-    guint i, n;
+    guint n;
 
     buf = g_ptr_array_index (pool->buffers, pool->current_buffer_index);
     g_assert (pool->other_pool == buf->pool);
     gst_object_replace ((GstObject **) & buf->pool, NULL);
 
     n = gst_buffer_n_memory (buf);
-    for (i = 0; i < n; i++) {
-      GstMemory *mem = gst_buffer_peek_memory (buf, i);
+    g_return_val_if_fail (n == 1, GST_FLOW_ERROR);
 
-      /* FIXME: We don't allow sharing because we need to know
-       * when the memory becomes unused and can only then put
-       * it back to the pool. Which is done in the pool's release
-       * function
-       */
-      GST_MINI_OBJECT_FLAG_SET (mem, GST_MEMORY_FLAG_NO_SHARE);
-    }
+    /* rip the memory out of the buffer;
+     * we like to keep them separate in this pool */
+    foreign_mem = gst_buffer_get_memory (buf, 0);
+    gst_buffer_remove_all_memory (buf);
 
     if (pool->add_videometa) {
       GstVideoMeta *meta;
@@ -440,41 +339,12 @@ gst_omx_buffer_pool_alloc_buffer (GstBufferPool * bpool,
 
     pool->need_copy = FALSE;
   } else {
-    GstMemory *mem;
     const guint nstride = pool->port->port_def.format.video.nStride;
     const guint nslice = pool->port->port_def.format.video.nSliceHeight;
     gsize offset[GST_VIDEO_MAX_PLANES] = { 0, };
     gint stride[GST_VIDEO_MAX_PLANES] = { nstride, 0, };
 
-    if (pool->output_mode == GST_OMX_BUFFER_MODE_DMABUF) {
-      gint fd;
-      GstMapInfo map;
-
-      fd = GPOINTER_TO_INT (omx_buf->omx_buf->pBuffer);
-
-      mem =
-          gst_dmabuf_allocator_alloc (pool->allocator, fd,
-          omx_buf->omx_buf->nAllocLen);
-
-      if (!gst_caps_features_contains (gst_caps_get_features (pool->caps, 0),
-              GST_CAPS_FEATURE_MEMORY_DMABUF)) {
-        /* Check if the memory is actually mappable */
-        if (!gst_memory_map (mem, &map, GST_MAP_READWRITE)) {
-          GST_ERROR_OBJECT (pool,
-              "dmabuf memory is not mappable but caps does not have the 'memory:DMABuf' feature");
-          gst_memory_unref (mem);
-          return GST_FLOW_ERROR;
-        }
-
-        gst_memory_unmap (mem, &map);
-      }
-    } else {
-      mem = gst_omx_memory_allocator_alloc (pool->allocator, 0, omx_buf);
-    }
-
     buf = gst_buffer_new ();
-    gst_buffer_append_memory (buf, mem);
-    g_ptr_array_add (pool->buffers, buf);
 
     switch (GST_VIDEO_INFO_FORMAT (&pool->video_info)) {
       case GST_VIDEO_FORMAT_ABGR:
@@ -545,7 +415,29 @@ gst_omx_buffer_pool_alloc_buffer (GstBufferPool * bpool,
     }
   }
 
-  gst_omx_buffer_set_omx_buf (buf, omx_buf);
+  mem = gst_omx_allocator_allocate (pool->allocator, pool->current_buffer_index,
+      foreign_mem);
+  if (!mem)
+    return GST_FLOW_ERROR;
+
+  if (pool->output_mode == GST_OMX_BUFFER_MODE_DMABUF) {
+    GstMapInfo map;
+
+    if (!gst_caps_features_contains (gst_caps_get_features (pool->caps, 0),
+            GST_CAPS_FEATURE_MEMORY_DMABUF)) {
+      /* Check if the memory is actually mappable */
+      if (!gst_memory_map (mem, &map, GST_MAP_READWRITE)) {
+        GST_ERROR_OBJECT (pool,
+            "dmabuf memory is not mappable but caps does not have the 'memory:DMABuf' feature");
+        gst_memory_unref (mem);
+        return GST_FLOW_ERROR;
+      }
+
+      gst_memory_unmap (mem, &map);
+    }
+  }
+
+  /* mem still belongs to the allocator; do not add it in the buffer just yet */
 
   *buffer = buf;
 
@@ -554,6 +446,22 @@ gst_omx_buffer_pool_alloc_buffer (GstBufferPool * bpool,
   return GST_FLOW_OK;
 }
 
+/* called by the allocator when we are using other_pool in order
+ * to restore the foreign GstMemory back to its original GstBuffer */
+static void
+on_allocator_foreign_mem_released (GstOMXAllocator * allocator,
+    gint index, GstMemory * mem, GstOMXBufferPool * pool)
+{
+  GstBuffer *buf;
+
+  buf = g_ptr_array_index (pool->buffers, index);
+  gst_buffer_append_memory (buf, mem);
+
+  /* the buffer consumed the passed reference.
+   * we still need one more reference for the allocator */
+  gst_memory_ref (mem);
+}
+
 static void
 gst_omx_buffer_pool_free_buffer (GstBufferPool * bpool, GstBuffer * buffer)
 {
@@ -567,59 +475,29 @@ gst_omx_buffer_pool_free_buffer (GstBufferPool * bpool, GstBuffer * buffer)
   }
   GST_OBJECT_UNLOCK (pool);
 
-  gst_omx_buffer_set_omx_buf (buffer, NULL);
-
   GST_BUFFER_POOL_CLASS (gst_omx_buffer_pool_parent_class)->free_buffer (bpool,
       buffer);
 }
 
-static GstBuffer *
-find_buffer_from_omx_buffer (GstOMXBufferPool * pool, GstOMXBuffer * omx_buf)
-{
-  guint i;
-
-  for (i = 0; i < pool->buffers->len; i++) {
-    GstBuffer *buf = g_ptr_array_index (pool->buffers, i);
-
-    if (gst_omx_buffer_get_omx_buf (buf) == omx_buf)
-      return buf;
-  }
-
-  return NULL;
-}
-
 static GstFlowReturn
 gst_omx_buffer_pool_acquire_buffer (GstBufferPool * bpool,
     GstBuffer ** buffer, GstBufferPoolAcquireParams * params)
 {
   GstFlowReturn ret;
   GstOMXBufferPool *pool = GST_OMX_BUFFER_POOL (bpool);
+  GstMemory *mem;
 
   if (pool->port->port_def.eDir == OMX_DirOutput) {
-    GstBuffer *buf;
-
     g_return_val_if_fail (pool->current_buffer_index != -1, GST_FLOW_ERROR);
 
-    buf = g_ptr_array_index (pool->buffers, pool->current_buffer_index);
-    g_return_val_if_fail (buf != NULL, GST_FLOW_ERROR);
-    *buffer = buf;
-    ret = GST_FLOW_OK;
+    ret = gst_omx_allocator_acquire (pool->allocator, &mem,
+        pool->current_buffer_index, NULL);
+    if (ret != GST_FLOW_OK)
+      return ret;
 
     /* If it's our own memory we have to set the sizes */
     if (!pool->other_pool) {
-      GstMemory *mem = gst_buffer_peek_memory (*buffer, 0);
-      GstOMXBuffer *omx_buf;
-
-      if (pool->output_mode == GST_OMX_BUFFER_MODE_DMABUF) {
-        omx_buf = gst_omx_buffer_get_omx_buf (buf);
-      } else {
-        g_assert (mem
-            && g_strcmp0 (mem->allocator->mem_type, GST_OMX_MEMORY_TYPE) == 0);
-        /* We already have a pointer to the GstOMXBuffer, no need to retrieve it
-         * from the qdata */
-        omx_buf = ((GstOMXMemory *) mem)->buf;
-      }
-
+      GstOMXBuffer *omx_buf = gst_omx_memory_get_omx_buf (mem);
       mem->size = omx_buf->omx_buf->nFilledLen;
       mem->offset = omx_buf->omx_buf->nOffset;
     }
@@ -634,9 +512,9 @@ gst_omx_buffer_pool_acquire_buffer (GstBufferPool * bpool,
 
     r = gst_omx_port_acquire_buffer (pool->port, &omx_buf, wait);
     if (r == GST_OMX_ACQUIRE_BUFFER_OK) {
-      *buffer = find_buffer_from_omx_buffer (pool, omx_buf);
-      g_return_val_if_fail (*buffer, GST_FLOW_ERROR);
-      return GST_FLOW_OK;
+      ret = gst_omx_allocator_acquire (pool->allocator, &mem, -1, omx_buf);
+      if (ret != GST_FLOW_OK)
+        return ret;
     } else if (r == GST_OMX_ACQUIRE_BUFFER_FLUSHING) {
       return GST_FLOW_FLUSHING;
     } else {
@@ -644,32 +522,65 @@ gst_omx_buffer_pool_acquire_buffer (GstBufferPool * bpool,
     }
   }
 
+  /* get some GstBuffer available in this pool */
+  ret = GST_BUFFER_POOL_CLASS (gst_omx_buffer_pool_parent_class)->acquire_buffer
+      (bpool, buffer, params);
+
+  if (ret == GST_FLOW_OK) {
+    /* attach the acquired memory on it */
+    gst_buffer_append_memory (*buffer, mem);
+  } else {
+    gst_memory_unref (mem);
+  }
+
   return ret;
 }
 
 static void
-gst_omx_buffer_pool_release_buffer (GstBufferPool * bpool, GstBuffer * buffer)
+gst_omx_buffer_pool_reset_buffer (GstBufferPool * bpool, GstBuffer * buffer)
 {
   GstOMXBufferPool *pool = GST_OMX_BUFFER_POOL (bpool);
+  guint n;
+
+  n = gst_buffer_n_memory (buffer);
+  if (G_UNLIKELY (n != 1)) {
+    GST_ERROR_OBJECT (pool, "Released buffer does not have 1 memory... "
+        "(n = %u) something went terribly wrong", n);
+  }
+
+  /* rip the memory out of the buffer;
+   * we like to keep them separate in this pool.
+   * if this was the last ref count of the memory, it will be returned
+   * to the allocator, otherwise it will be returned later */
+  gst_buffer_remove_all_memory (buffer);
+
+  /* reset before removing the TAG_MEMORY flag so that the parent impl
+   * doesn't try to restore the original buffer size */
+  GST_BUFFER_POOL_CLASS (gst_omx_buffer_pool_parent_class)->reset_buffer
+      (bpool, buffer);
+
+  /* pretend nothing happened to the memory to avoid discarding the buffer */
+  GST_MINI_OBJECT_FLAG_UNSET (buffer, GST_BUFFER_FLAG_TAG_MEMORY);
+}
+
+static void
+on_allocator_omxbuf_released (GstOMXAllocator * allocator,
+    GstOMXBuffer * omx_buf, GstOMXBufferPool * pool)
+{
   OMX_ERRORTYPE err;
-  GstOMXBuffer *omx_buf;
-
-  g_assert (pool->component && pool->port);
-
-  if (gst_buffer_pool_is_active (bpool)) {
-    omx_buf = gst_omx_buffer_get_omx_buf (buffer);
-    if (pool->port->port_def.eDir == OMX_DirOutput && !omx_buf->used &&
-        !pool->deactivated) {
-      /* Release back to the port, can be filled again */
-      err = gst_omx_port_release_buffer (pool->port, omx_buf);
-      if (err != OMX_ErrorNone) {
-        GST_ELEMENT_ERROR (pool->element, LIBRARY, SETTINGS, (NULL),
-            ("Failed to relase output buffer to component: %s (0x%08x)",
-                gst_omx_error_to_string (err), err));
-      }
-    } else if (pool->port->port_def.eDir == OMX_DirInput) {
-      gst_omx_port_requeue_buffer (pool->port, omx_buf);
+
+  if (pool->port->port_def.eDir == OMX_DirOutput && !omx_buf->used &&
+      !pool->deactivated) {
+    /* Release back to the port, can be filled again */
+    err = gst_omx_port_release_buffer (pool->port, omx_buf);
+
+    if (err != OMX_ErrorNone) {
+      GST_ELEMENT_ERROR (pool->element, LIBRARY, SETTINGS, (NULL),
+          ("Failed to relase output buffer to component: %s (0x%08x)",
+              gst_omx_error_to_string (err), err));
     }
+  } else if (pool->port->port_def.eDir == OMX_DirInput) {
+    gst_omx_port_requeue_buffer (pool->port, omx_buf);
   }
 }
 
@@ -717,7 +628,7 @@ gst_omx_buffer_pool_class_init (GstOMXBufferPoolClass * klass)
   gstbufferpool_class->alloc_buffer = gst_omx_buffer_pool_alloc_buffer;
   gstbufferpool_class->free_buffer = gst_omx_buffer_pool_free_buffer;
   gstbufferpool_class->acquire_buffer = gst_omx_buffer_pool_acquire_buffer;
-  gstbufferpool_class->release_buffer = gst_omx_buffer_pool_release_buffer;
+  gstbufferpool_class->reset_buffer = gst_omx_buffer_pool_reset_buffer;
 
   signals[SIG_ALLOCATE] = g_signal_new ("allocate",
       G_TYPE_FROM_CLASS (gobject_class), G_SIGNAL_RUN_LAST, 0, NULL, NULL, NULL,
@@ -741,18 +652,12 @@ gst_omx_buffer_pool_new (GstElement * element, GstOMXComponent * component,
   pool->component = gst_omx_component_ref (component);
   pool->port = port;
   pool->output_mode = output_mode;
+  pool->allocator = gst_omx_allocator_new (component, port);
 
-  switch (output_mode) {
-    case GST_OMX_BUFFER_MODE_DMABUF:
-      pool->allocator = gst_dmabuf_allocator_new ();
-      break;
-    case GST_OMX_BUFFER_MODE_SYSTEM_MEMORY:
-      pool->allocator =
-          g_object_new (gst_omx_memory_allocator_get_type (), NULL);
-      break;
-    default:
-      g_assert_not_reached ();
-  }
+  g_signal_connect_object (pool->allocator, "omxbuf-released",
+      (GCallback) on_allocator_omxbuf_released, pool, 0);
+  g_signal_connect_object (pool->allocator, "foreign-mem-released",
+      (GCallback) on_allocator_foreign_mem_released, pool, 0);
 
   return GST_BUFFER_POOL (pool);
 }
diff --git a/omx/gstomxbufferpool.h b/omx/gstomxbufferpool.h
index a09c825..bc5ac60 100644
--- a/omx/gstomxbufferpool.h
+++ b/omx/gstomxbufferpool.h
@@ -30,6 +30,7 @@
 #include <gst/video/gstvideopool.h>
 
 #include "gstomx.h"
+#include "gstomxallocator.h"
 
 G_BEGIN_DECLS
 
@@ -65,7 +66,7 @@ struct _GstOMXBufferPool
   GstOMXPort *port;
 
   /* For handling OpenMAX allocated memory */
-  GstAllocator *allocator;
+  GstOMXAllocator *allocator;
 
   /* Set from outside this pool */
   /* TRUE if the pool is not used anymore */
diff --git a/omx/gstomxvideoenc.c b/omx/gstomxvideoenc.c
index dfc03c5..7e088fc 100644
--- a/omx/gstomxvideoenc.c
+++ b/omx/gstomxvideoenc.c
@@ -2040,6 +2040,15 @@ gst_omx_video_enc_set_to_idle (GstOMXVideoEnc * self)
   return TRUE;
 }
 
+static GstOMXBuffer *
+get_omx_buf (GstBuffer * buffer)
+{
+  GstMemory *mem;
+
+  mem = gst_buffer_peek_memory (buffer, 0);
+  return gst_omx_memory_get_omx_buf (mem);
+}
+
 static gboolean
 buffer_is_from_input_pool (GstOMXVideoEnc * self, GstBuffer * buffer)
 {
@@ -2047,7 +2056,7 @@ buffer_is_from_input_pool (GstOMXVideoEnc * self, GstBuffer * buffer)
    * with our input port. */
   GstOMXBuffer *buf;
 
-  buf = gst_omx_buffer_get_omx_buf (buffer);
+  buf = get_omx_buf (buffer);
   if (!buf)
     return FALSE;
 
@@ -2761,7 +2770,7 @@ gst_omx_video_enc_handle_frame (GstVideoEncoder * encoder,
 
     if (buffer_is_from_input_pool (self, frame->input_buffer)) {
       /* Receiving a buffer from our input pool */
-      buf = gst_omx_buffer_get_omx_buf (frame->input_buffer);
+      buf = get_omx_buf (frame->input_buffer);
 
       GST_LOG_OBJECT (self,
           "Input buffer %p already has a OMX buffer associated: %p",
diff --git a/omx/meson.build b/omx/meson.build
index b68cc55..1bcebf5 100644
--- a/omx/meson.build
+++ b/omx/meson.build
@@ -1,5 +1,6 @@
 omx_sources = [
   'gstomx.c',
+  'gstomxallocator.c',
   'gstomxbufferpool.c',
   'gstomxvideo.c',
   'gstomxvideodec.c',