#include "q_struct.h"
 #include "nicvf_queues.h"
 
+static void nicvf_get_page(struct nicvf *nic)
+{
+       if (!nic->rb_pageref || !nic->rb_page)
+               return;
+
+       atomic_add(nic->rb_pageref, &nic->rb_page->_count);
+       nic->rb_pageref = 0;
+}
+
 /* Poll a register for a specific value */
 static int nicvf_poll_reg(struct nicvf *nic, int qidx,
                          u64 reg, int bit_pos, int bits, int val)
        int order = (PAGE_SIZE <= 4096) ?  PAGE_ALLOC_COSTLY_ORDER : 0;
 
        /* Check if request can be accomodated in previous allocated page */
-       if (nic->rb_page) {
-               if ((nic->rb_page_offset + buf_len + buf_len) >
-                   (PAGE_SIZE << order)) {
-                       nic->rb_page = NULL;
-               } else {
-                       nic->rb_page_offset += buf_len;
-                       get_page(nic->rb_page);
-               }
+       if (nic->rb_page &&
+           ((nic->rb_page_offset + buf_len) < (PAGE_SIZE << order))) {
+               nic->rb_pageref++;
+               goto ret;
        }
 
+       nicvf_get_page(nic);
+       nic->rb_page = NULL;
+
        /* Allocate a new page */
        if (!nic->rb_page) {
                nic->rb_page = alloc_pages(gfp | __GFP_COMP | __GFP_NOWARN,
                nic->rb_page_offset = 0;
        }
 
+ret:
        *rbuf = (u64 *)((u64)page_address(nic->rb_page) + nic->rb_page_offset);
+       nic->rb_page_offset += buf_len;
 
        return 0;
 }
                desc = GET_RBDR_DESC(rbdr, idx);
                desc->buf_addr = virt_to_phys(rbuf) >> NICVF_RCV_BUF_ALIGN;
        }
+
+       nicvf_get_page(nic);
+
        return 0;
 }
 
                new_rb++;
        }
 
+       nicvf_get_page(nic);
+
        /* make sure all memory stores are done before ringing doorbell */
        smp_wmb();