1 files changed, 15 insertions, 8 deletions

diff --git a/drivers/staging/media/atomisp/pci/runtime/rmgr/src/rmgr_vbuf.c b/drivers/staging/media/atomisp/pci/runtime/rmgr/src/rmgr_vbuf.c
index d96aaa4bc75d..2e07dab8bf51 100644
--- a/drivers/staging/media/atomisp/pci/runtime/rmgr/src/rmgr_vbuf.c
+++ b/drivers/staging/media/atomisp/pci/runtime/rmgr/src/rmgr_vbuf.c
@@ -254,14 +254,15 @@ void rmgr_pop_handle(struct ia_css_rmgr_vbuf_pool *pool,
 void ia_css_rmgr_acq_vbuf(struct ia_css_rmgr_vbuf_pool *pool,
 			  struct ia_css_rmgr_vbuf_handle **handle)
 {
-	struct ia_css_rmgr_vbuf_handle h;
-
 	if ((!pool) || (!handle) || (!*handle)) {
 		IA_CSS_LOG("Invalid inputs");
 		return;
 	}
 
 	if (pool->copy_on_write) {
+		struct ia_css_rmgr_vbuf_handle *new_handle;
+		struct ia_css_rmgr_vbuf_handle h = { 0 };
+
 		/* only one reference, reuse (no new retain) */
 		if ((*handle)->count == 1)
 			return;
@@ -272,23 +273,29 @@ void ia_css_rmgr_acq_vbuf(struct ia_css_rmgr_vbuf_pool *pool,
 			h.size = (*handle)->size;
 			/* release ref to current buffer */
 			ia_css_rmgr_refcount_release_vbuf(handle);
-			*handle = &h;
+			new_handle = &h;
+		} else {
+			new_handle = *handle;
 		}
 		/* get new buffer for needed size */
-		if ((*handle)->vptr == 0x0) {
+		if (new_handle->vptr == 0x0) {
 			if (pool->recycle) {
 				/* try and pop from pool */
-				rmgr_pop_handle(pool, handle);
+				rmgr_pop_handle(pool, &new_handle);
 			}
-			if ((*handle)->vptr == 0x0) {
+			if (new_handle->vptr == 0x0) {
 				/* we need to allocate */
-				(*handle)->vptr = hmm_alloc((*handle)->size,
-							     HMM_BO_PRIVATE, 0, NULL, 0);
+				new_handle->vptr = hmm_alloc(new_handle->size);
 			} else {
 				/* we popped a buffer */
+				*handle = new_handle;
 				return;
 			}
 		}
+		/* Note that new_handle will change to an internally maintained one */
+		ia_css_rmgr_refcount_retain_vbuf(&new_handle);
+		*handle = new_handle;
+		return;
 	}
 	/* Note that handle will change to an internally maintained one */
 	ia_css_rmgr_refcount_retain_vbuf(handle);