[PATCH v2 0/3] mm: kmemleak: Improve vmalloc() false positives for thread stack allocation

[PATCH v2 0/3] mm: kmemleak: Improve vmalloc() false positives for thread stack allocation

[Catalin Marinas]

Hi,

This is a follow up from [1] (mm: kmemleak: Treat vm_struct as
alternative reference to vmalloc'ed objects).

The first two patches are just clean-up and refactoring. The third
introduces the kmemleak_vmalloc() API which allows a vmalloc() caller to
keep either the returned pointer or a pointer to vm_struct as a
reference (see the patch description for the implementation details).
The false positives were noticed with alloc_thread_stack_node(),
free_thread_stack() and CONFIG_VMAP_STACK where a per-CPU array is used
to cache the freed thread stacks as vm_struct pointers.

Changes since v1:

- Split the patch into three for easier review
- Only call update_refs() if !color_gray() on the found object, it
  avoids an unnecessary function call

[1] http://lkml.kernel.org/r/1495474514-24425-1-git-send-email-catalin.marinas@arm.com

Catalin Marinas (3):
  mm: kmemleak: Slightly reduce the size of some structures on 64-bit
    architectures
  mm: kmemleak: Factor object reference updating out of scan_block()
  mm: kmemleak: Treat vm_struct as alternative reference to vmalloc'ed
    objects

 Documentation/dev-tools/kmemleak.rst |   1 +
 include/linux/kmemleak.h             |   7 ++
 mm/kmemleak.c                        | 136 +++++++++++++++++++++++++++++------
 mm/vmalloc.c                         |   7 +-
 4 files changed, 123 insertions(+), 28 deletions(-)

[PATCH v2 1/3] mm: kmemleak: Slightly reduce the size of some structures on 64-bit architectures

[Catalin Marinas]

This patch changes the kmemleak_object.flags type to unsigned int and
moves the early_log.min_count (int) near early_log.op_type (int) to
slightly reduce the size of these structures on 64-bit architectures.

Cc: Michal Hocko <mhocko@kernel.org>
Cc: Andy Lutomirski <luto@amacapital.net>
Cc: "Luis R. Rodriguez" <mcgrof@kernel.org>
Signed-off-by: Catalin Marinas <catalin.marinas@arm.com>
---
 mm/kmemleak.c | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/mm/kmemleak.c b/mm/kmemleak.c
index 20036d4f9f13..964b12eba2c1 100644
--- a/mm/kmemleak.c
+++ b/mm/kmemleak.c
@@ -150,7 +150,7 @@ struct kmemleak_scan_area {
  */
 struct kmemleak_object {
 	spinlock_t lock;
-	unsigned long flags;		/* object status flags */
+	unsigned int flags;		/* object status flags */
 	struct list_head object_list;
 	struct list_head gray_list;
 	struct rb_node rb_node;
@@ -262,9 +262,9 @@ enum {
  */
 struct early_log {
 	int op_type;			/* kmemleak operation type */
+	int min_count;			/* minimum reference count */
 	const void *ptr;		/* allocated/freed memory block */
 	size_t size;			/* memory block size */
-	int min_count;			/* minimum reference count */
 	unsigned long trace[MAX_TRACE];	/* stack trace */
 	unsigned int trace_len;		/* stack trace length */
 };
@@ -393,7 +393,7 @@ static void dump_object_info(struct kmemleak_object *object)
 		  object->comm, object->pid, object->jiffies);
 	pr_notice("  min_count = %d\n", object->min_count);
 	pr_notice("  count = %d\n", object->count);
-	pr_notice("  flags = 0x%lx\n", object->flags);
+	pr_notice("  flags = 0x%x\n", object->flags);
 	pr_notice("  checksum = %u\n", object->checksum);
 	pr_notice("  backtrace:\n");
 	print_stack_trace(&trace, 4);

[PATCH v2 2/3] mm: kmemleak: Factor object reference updating out of scan_block()

[Catalin Marinas]

The scan_block() function updates the number of references (pointers) to
objects, adding them to the gray_list when object->min_count is reached.
The patch factors out this functionality into a separate update_refs()
function.

Cc: Michal Hocko <mhocko@kernel.org>
Cc: Andy Lutomirski <luto@amacapital.net>
Cc: "Luis R. Rodriguez" <mcgrof@kernel.org>
Signed-off-by: Catalin Marinas <catalin.marinas@arm.com>
---
 mm/kmemleak.c | 43 +++++++++++++++++++++++++------------------
 1 file changed, 25 insertions(+), 18 deletions(-)

diff --git a/mm/kmemleak.c b/mm/kmemleak.c
index 964b12eba2c1..266482f460c2 100644
--- a/mm/kmemleak.c
+++ b/mm/kmemleak.c
@@ -1188,6 +1188,30 @@ static bool update_checksum(struct kmemleak_object *object)
 }
 
 /*
+ * Update an object's references. object->lock must be held by the caller.
+ */
+static void update_refs(struct kmemleak_object *object)
+{
+	if (!color_white(object)) {
+		/* non-orphan, ignored or new */
+		return;
+	}
+
+	/*
+	 * Increase the object's reference count (number of pointers to the
+	 * memory block). If this count reaches the required minimum, the
+	 * object's color will become gray and it will be added to the
+	 * gray_list.
+	 */
+	object->count++;
+	if (color_gray(object)) {
+		/* put_object() called when removing from gray_list */
+		WARN_ON(!get_object(object));
+		list_add_tail(&object->gray_list, &gray_list);
+	}
+}
+
+/*
  * Memory scanning is a long process and it needs to be interruptable. This
  * function checks whether such interrupt condition occurred.
  */
@@ -1259,24 +1283,7 @@ static void scan_block(void *_start, void *_end,
 		 * enclosed by scan_mutex.
 		 */
 		spin_lock_nested(&object->lock, SINGLE_DEPTH_NESTING);
-		if (!color_white(object)) {
-			/* non-orphan, ignored or new */
-			spin_unlock(&object->lock);
-			continue;
-		}
-
-		/*
-		 * Increase the object's reference count (number of pointers
-		 * to the memory block). If this count reaches the required
-		 * minimum, the object's color will become gray and it will be
-		 * added to the gray_list.
-		 */
-		object->count++;
-		if (color_gray(object)) {
-			/* put_object() called when removing from gray_list */
-			WARN_ON(!get_object(object));
-			list_add_tail(&object->gray_list, &gray_list);
-		}
+		update_refs(object);
 		spin_unlock(&object->lock);
 	}
 	read_unlock_irqrestore(&kmemleak_lock, flags);

Re: [PATCH v2 2/3] mm: kmemleak: Factor object reference updating out of scan_block()

[Luis Henriques]

On Thu, May 25, 2017 at 04:42:16PM +0100, Catalin Marinas wrote:
> The scan_block() function updates the number of references (pointers) to
> objects, adding them to the gray_list when object->min_count is reached.
> The patch factors out this functionality into a separate update_refs()
> function.
> 
> Cc: Michal Hocko <mhocko@kernel.org>
> Cc: Andy Lutomirski <luto@amacapital.net>
> Cc: "Luis R. Rodriguez" <mcgrof@kernel.org>
> Signed-off-by: Catalin Marinas <catalin.marinas@arm.com>
> ---
>  mm/kmemleak.c | 43 +++++++++++++++++++++++++------------------
>  1 file changed, 25 insertions(+), 18 deletions(-)
> 
> diff --git a/mm/kmemleak.c b/mm/kmemleak.c
> index 964b12eba2c1..266482f460c2 100644
> --- a/mm/kmemleak.c
> +++ b/mm/kmemleak.c
> @@ -1188,6 +1188,30 @@ static bool update_checksum(struct kmemleak_object *object)
>  }
>  
>  /*
> + * Update an object's references. object->lock must be held by the caller.
> + */
> +static void update_refs(struct kmemleak_object *object)
> +{
> +	if (!color_white(object)) {
> +		/* non-orphan, ignored or new */
> +		return;
> +	}
> +
> +	/*
> +	 * Increase the object's reference count (number of pointers to the
> +	 * memory block). If this count reaches the required minimum, the
> +	 * object's color will become gray and it will be added to the
> +	 * gray_list.
> +	 */
> +	object->count++;
> +	if (color_gray(object)) {
> +		/* put_object() called when removing from gray_list */
> +		WARN_ON(!get_object(object));
> +		list_add_tail(&object->gray_list, &gray_list);
> +	}
> +}
> +
> +/*
>   * Memory scanning is a long process and it needs to be interruptable. This
>   * function checks whether such interrupt condition occurred.
>   */
> @@ -1259,24 +1283,7 @@ static void scan_block(void *_start, void *_end,
>  		 * enclosed by scan_mutex.
>  		 */
>  		spin_lock_nested(&object->lock, SINGLE_DEPTH_NESTING);
> -		if (!color_white(object)) {
> -			/* non-orphan, ignored or new */
> -			spin_unlock(&object->lock);
> -			continue;
> -		}
> -
> -		/*
> -		 * Increase the object's reference count (number of pointers
> -		 * to the memory block). If this count reaches the required
> -		 * minimum, the object's color will become gray and it will be
> -		 * added to the gray_list.
> -		 */
> -		object->count++;
> -		if (color_gray(object)) {
> -			/* put_object() called when removing from gray_list */
> -			WARN_ON(!get_object(object));
> -			list_add_tail(&object->gray_list, &gray_list);
> -		}
> +		update_refs(object);
>  		spin_unlock(&object->lock);

FWIW, I've tested this patchset and I don't see kmemleak triggering the
false positives anymore.

I've also done a quick review and couldn't find anything obviously
incorrect, just a question: why didn't you moved the spin_lock/unlock into
update_refs() too?  It would save you 2 lines in the next patch :)

Cheers,
--
Luís


>  	}
>  	read_unlock_irqrestore(&kmemleak_lock, flags);
> 

Re: [PATCH v2 2/3] mm: kmemleak: Factor object reference updating out of scan_block()

[Catalin Marinas]

On Fri, May 26, 2017 at 05:09:17PM +0100, Luis Henriques wrote:
> On Thu, May 25, 2017 at 04:42:16PM +0100, Catalin Marinas wrote:
> > The scan_block() function updates the number of references (pointers) to
> > objects, adding them to the gray_list when object->min_count is reached.
> > The patch factors out this functionality into a separate update_refs()
> > function.
> > 
> > Cc: Michal Hocko <mhocko@kernel.org>
> > Cc: Andy Lutomirski <luto@amacapital.net>
> > Cc: "Luis R. Rodriguez" <mcgrof@kernel.org>
> > Signed-off-by: Catalin Marinas <catalin.marinas@arm.com>
> > ---
> >  mm/kmemleak.c | 43 +++++++++++++++++++++++++------------------
> >  1 file changed, 25 insertions(+), 18 deletions(-)
> > 
> > diff --git a/mm/kmemleak.c b/mm/kmemleak.c
> > index 964b12eba2c1..266482f460c2 100644
> > --- a/mm/kmemleak.c
> > +++ b/mm/kmemleak.c
> > @@ -1188,6 +1188,30 @@ static bool update_checksum(struct kmemleak_object *object)
> >  }
> >  
> >  /*
> > + * Update an object's references. object->lock must be held by the caller.
> > + */
> > +static void update_refs(struct kmemleak_object *object)
> > +{
> > +	if (!color_white(object)) {
> > +		/* non-orphan, ignored or new */
> > +		return;
> > +	}
> > +
> > +	/*
> > +	 * Increase the object's reference count (number of pointers to the
> > +	 * memory block). If this count reaches the required minimum, the
> > +	 * object's color will become gray and it will be added to the
> > +	 * gray_list.
> > +	 */
> > +	object->count++;
> > +	if (color_gray(object)) {
> > +		/* put_object() called when removing from gray_list */
> > +		WARN_ON(!get_object(object));
> > +		list_add_tail(&object->gray_list, &gray_list);
> > +	}
> > +}
> > +
> > +/*
> >   * Memory scanning is a long process and it needs to be interruptable. This
> >   * function checks whether such interrupt condition occurred.
> >   */
> > @@ -1259,24 +1283,7 @@ static void scan_block(void *_start, void *_end,
> >  		 * enclosed by scan_mutex.
> >  		 */
> >  		spin_lock_nested(&object->lock, SINGLE_DEPTH_NESTING);
> > -		if (!color_white(object)) {
> > -			/* non-orphan, ignored or new */
> > -			spin_unlock(&object->lock);
> > -			continue;
> > -		}
> > -
> > -		/*
> > -		 * Increase the object's reference count (number of pointers
> > -		 * to the memory block). If this count reaches the required
> > -		 * minimum, the object's color will become gray and it will be
> > -		 * added to the gray_list.
> > -		 */
> > -		object->count++;
> > -		if (color_gray(object)) {
> > -			/* put_object() called when removing from gray_list */
> > -			WARN_ON(!get_object(object));
> > -			list_add_tail(&object->gray_list, &gray_list);
> > -		}
> > +		update_refs(object);
> >  		spin_unlock(&object->lock);
> 
> FWIW, I've tested this patchset and I don't see kmemleak triggering the
> false positives anymore.

Thanks for re-testing (I dropped your tested-by from the initial patch
since I made a small modification).

> I've also done a quick review and couldn't find anything obviously
> incorrect, just a question: why didn't you moved the spin_lock/unlock into
> update_refs() too?  It would save you 2 lines in the next patch :)

There is a small difference: for the first object it needs to check
color_gray() and access object->excess_ref while the lock is held. It
doesn't need this in the second case. I could've written it in different
ways but probably with a similar number of lines; I just found this
clearer.

-- 
Catalin

Re: [PATCH v2 2/3] mm: kmemleak: Factor object reference updating out of scan_block()

[Catalin Marinas]

On Fri, May 26, 2017 at 05:21:08PM +0100, Catalin Marinas wrote:
> On Fri, May 26, 2017 at 05:09:17PM +0100, Luis Henriques wrote:
> > On Thu, May 25, 2017 at 04:42:16PM +0100, Catalin Marinas wrote:
> > > The scan_block() function updates the number of references (pointers) to
> > > objects, adding them to the gray_list when object->min_count is reached.
> > > The patch factors out this functionality into a separate update_refs()
> > > function.
> > > 
> > > Cc: Michal Hocko <mhocko@kernel.org>
> > > Cc: Andy Lutomirski <luto@amacapital.net>
> > > Cc: "Luis R. Rodriguez" <mcgrof@kernel.org>
> > > Signed-off-by: Catalin Marinas <catalin.marinas@arm.com>
[...]
> > FWIW, I've tested this patchset and I don't see kmemleak triggering the
> > false positives anymore.
> 
> Thanks for re-testing (I dropped your tested-by from the initial patch
> since I made a small modification).

Sorry, the "re-testing" comment was meant at the other Luis on cc ;)
(Luis R. Rodriguez). It's been a long day...

-- 
Catalin

Re: [PATCH v2 2/3] mm: kmemleak: Factor object reference updating out of scan_block()

[Luis Henriques]

On Fri, May 26, 2017 at 05:23:30PM +0100, Catalin Marinas wrote:
> On Fri, May 26, 2017 at 05:21:08PM +0100, Catalin Marinas wrote:
> > On Fri, May 26, 2017 at 05:09:17PM +0100, Luis Henriques wrote:
> > > On Thu, May 25, 2017 at 04:42:16PM +0100, Catalin Marinas wrote:
> > > > The scan_block() function updates the number of references (pointers) to
> > > > objects, adding them to the gray_list when object->min_count is reached.
> > > > The patch factors out this functionality into a separate update_refs()
> > > > function.
> > > > 
> > > > Cc: Michal Hocko <mhocko@kernel.org>
> > > > Cc: Andy Lutomirski <luto@amacapital.net>
> > > > Cc: "Luis R. Rodriguez" <mcgrof@kernel.org>
> > > > Signed-off-by: Catalin Marinas <catalin.marinas@arm.com>
> [...]
> > > FWIW, I've tested this patchset and I don't see kmemleak triggering the
> > > false positives anymore.
> > 
> > Thanks for re-testing (I dropped your tested-by from the initial patch
> > since I made a small modification).
> 
> Sorry, the "re-testing" comment was meant at the other Luis on cc ;)
> (Luis R. Rodriguez). It's been a long day...

Heh, no worries!  What are the odds of having 2 different guys named Luis
testing the same patch? :-)

Cheers,
--
Luís

[PATCH v2 3/3] mm: kmemleak: Treat vm_struct as alternative reference to vmalloc'ed objects

[Catalin Marinas]

Kmemleak requires that vmalloc'ed objects have a minimum reference count
of 2: one in the corresponding vm_struct object and the other owned by
the vmalloc() caller. There are cases, however, where the original
vmalloc() returned pointer is lost and, instead, a pointer to vm_struct
is stored (see free_thread_stack()). Kmemleak currently reports such
objects as leaks.

This patch adds support for treating any surplus references to an object
as additional references to a specified object. It introduces the
kmemleak_vmalloc() API function which takes a vm_struct pointer and sets
its surplus reference passing to the actual vmalloc() returned pointer.
The __vmalloc_node_range() calling site has been modified accordingly.

Cc: Michal Hocko <mhocko@kernel.org>
Cc: Andy Lutomirski <luto@amacapital.net>
Cc: "Luis R. Rodriguez" <mcgrof@kernel.org>
Reported-by: "Luis R. Rodriguez" <mcgrof@kernel.org>
Signed-off-by: Catalin Marinas <catalin.marinas@arm.com>
---
 Documentation/dev-tools/kmemleak.rst |  1 +
 include/linux/kmemleak.h             |  7 +++
 mm/kmemleak.c                        | 93 ++++++++++++++++++++++++++++++++++--
 mm/vmalloc.c                         |  7 +--
 4 files changed, 98 insertions(+), 10 deletions(-)

diff --git a/Documentation/dev-tools/kmemleak.rst b/Documentation/dev-tools/kmemleak.rst
index b2391b829169..cb8862659178 100644
--- a/Documentation/dev-tools/kmemleak.rst
+++ b/Documentation/dev-tools/kmemleak.rst
@@ -150,6 +150,7 @@ See the include/linux/kmemleak.h header for the functions prototype.
 - ``kmemleak_init``		 - initialize kmemleak
 - ``kmemleak_alloc``		 - notify of a memory block allocation
 - ``kmemleak_alloc_percpu``	 - notify of a percpu memory block allocation
+- ``kmemleak_vmalloc``		 - notify of a vmalloc() memory allocation
 - ``kmemleak_free``		 - notify of a memory block freeing
 - ``kmemleak_free_part``	 - notify of a partial memory block freeing
 - ``kmemleak_free_percpu``	 - notify of a percpu memory block freeing
diff --git a/include/linux/kmemleak.h b/include/linux/kmemleak.h
index 1c2a32829620..590343f6c1b1 100644
--- a/include/linux/kmemleak.h
+++ b/include/linux/kmemleak.h
@@ -22,6 +22,7 @@
 #define __KMEMLEAK_H
 
 #include <linux/slab.h>
+#include <linux/vmalloc.h>
 
 #ifdef CONFIG_DEBUG_KMEMLEAK
 
@@ -30,6 +31,8 @@ extern void kmemleak_alloc(const void *ptr, size_t size, int min_count,
 			   gfp_t gfp) __ref;
 extern void kmemleak_alloc_percpu(const void __percpu *ptr, size_t size,
 				  gfp_t gfp) __ref;
+extern void kmemleak_vmalloc(const struct vm_struct *area, size_t size,
+			     gfp_t gfp) __ref;
 extern void kmemleak_free(const void *ptr) __ref;
 extern void kmemleak_free_part(const void *ptr, size_t size) __ref;
 extern void kmemleak_free_percpu(const void __percpu *ptr) __ref;
@@ -81,6 +84,10 @@ static inline void kmemleak_alloc_percpu(const void __percpu *ptr, size_t size,
 					 gfp_t gfp)
 {
 }
+static inline void kmemleak_vmalloc(const struct vm_struct *area, size_t size,
+				    gfp_t gfp)
+{
+}
 static inline void kmemleak_free(const void *ptr)
 {
 }
diff --git a/mm/kmemleak.c b/mm/kmemleak.c
index 266482f460c2..7780cd83a495 100644
--- a/mm/kmemleak.c
+++ b/mm/kmemleak.c
@@ -159,6 +159,8 @@ struct kmemleak_object {
 	atomic_t use_count;
 	unsigned long pointer;
 	size_t size;
+	/* pass surplus references to this pointer */
+	unsigned long excess_ref;
 	/* minimum number of a pointers found before it is considered leak */
 	int min_count;
 	/* the total number of pointers found pointing to this object */
@@ -253,7 +255,8 @@ enum {
 	KMEMLEAK_NOT_LEAK,
 	KMEMLEAK_IGNORE,
 	KMEMLEAK_SCAN_AREA,
-	KMEMLEAK_NO_SCAN
+	KMEMLEAK_NO_SCAN,
+	KMEMLEAK_SET_EXCESS_REF
 };
 
 /*
@@ -264,7 +267,10 @@ struct early_log {
 	int op_type;			/* kmemleak operation type */
 	int min_count;			/* minimum reference count */
 	const void *ptr;		/* allocated/freed memory block */
-	size_t size;			/* memory block size */
+	union {
+		size_t size;		/* memory block size */
+		unsigned long excess_ref; /* surplus reference passing */
+	};
 	unsigned long trace[MAX_TRACE];	/* stack trace */
 	unsigned int trace_len;		/* stack trace length */
 };
@@ -562,6 +568,7 @@ static struct kmemleak_object *create_object(unsigned long ptr, size_t size,
 	object->flags = OBJECT_ALLOCATED;
 	object->pointer = ptr;
 	object->size = size;
+	object->excess_ref = 0;
 	object->min_count = min_count;
 	object->count = 0;			/* white color initially */
 	object->jiffies = jiffies;
@@ -795,6 +802,30 @@ static void add_scan_area(unsigned long ptr, size_t size, gfp_t gfp)
 }
 
 /*
+ * Any surplus references (object already gray) to 'ptr' are passed to
+ * 'excess_ref'. This is used in the vmalloc() case where a pointer to
+ * vm_struct may be used as an alternative reference to the vmalloc'ed object
+ * (see free_thread_stack()).
+ */
+static void object_set_excess_ref(unsigned long ptr, unsigned long excess_ref)
+{
+	unsigned long flags;
+	struct kmemleak_object *object;
+
+	object = find_and_get_object(ptr, 0);
+	if (!object) {
+		kmemleak_warn("Setting excess_ref on unknown object at 0x%08lx\n",
+			      ptr);
+		return;
+	}
+
+	spin_lock_irqsave(&object->lock, flags);
+	object->excess_ref = excess_ref;
+	spin_unlock_irqrestore(&object->lock, flags);
+	put_object(object);
+}
+
+/*
  * Set the OBJECT_NO_SCAN flag for the object corresponding to the give
  * pointer. Such object will not be scanned by kmemleak but references to it
  * are searched.
@@ -908,7 +939,7 @@ static void early_alloc_percpu(struct early_log *log)
  * @gfp:	kmalloc() flags used for kmemleak internal memory allocations
  *
  * This function is called from the kernel allocators when a new object
- * (memory block) is allocated (kmem_cache_alloc, kmalloc, vmalloc etc.).
+ * (memory block) is allocated (kmem_cache_alloc, kmalloc etc.).
  */
 void __ref kmemleak_alloc(const void *ptr, size_t size, int min_count,
 			  gfp_t gfp)
@@ -952,6 +983,36 @@ void __ref kmemleak_alloc_percpu(const void __percpu *ptr, size_t size,
 EXPORT_SYMBOL_GPL(kmemleak_alloc_percpu);
 
 /**
+ * kmemleak_vmalloc - register a newly vmalloc'ed object
+ * @area:	pointer to vm_struct
+ * @size:	size of the object
+ * @gfp:	__vmalloc() flags used for kmemleak internal memory allocations
+ *
+ * This function is called from the vmalloc() kernel allocator when a new
+ * object (memory block) is allocated.
+ */
+void __ref kmemleak_vmalloc(const struct vm_struct *area, size_t size, gfp_t gfp)
+{
+	pr_debug("%s(0x%p, %zu)\n", __func__, area, size);
+
+	/*
+	 * A min_count = 2 is needed because vm_struct contains a reference to
+	 * the virtual address of the vmalloc'ed block.
+	 */
+	if (kmemleak_enabled) {
+		create_object((unsigned long)area->addr, size, 2, gfp);
+		object_set_excess_ref((unsigned long)area,
+				      (unsigned long)area->addr);
+	} else if (kmemleak_early_log) {
+		log_early(KMEMLEAK_ALLOC, area->addr, size, 2);
+		/* reusing early_log.size for storing area->addr */
+		log_early(KMEMLEAK_SET_EXCESS_REF,
+			  area, (unsigned long)area->addr, 0);
+	}
+}
+EXPORT_SYMBOL_GPL(kmemleak_vmalloc);
+
+/**
  * kmemleak_free - unregister a previously registered object
  * @ptr:	pointer to beginning of the object
  *
@@ -1248,6 +1309,7 @@ static void scan_block(void *_start, void *_end,
 	for (ptr = start; ptr < end; ptr++) {
 		struct kmemleak_object *object;
 		unsigned long pointer;
+		unsigned long excess_ref;
 
 		if (scan_should_stop())
 			break;
@@ -1283,8 +1345,27 @@ static void scan_block(void *_start, void *_end,
 		 * enclosed by scan_mutex.
 		 */
 		spin_lock_nested(&object->lock, SINGLE_DEPTH_NESTING);
-		update_refs(object);
+		/* only pass surplus references (object already gray) */
+		if (color_gray(object)) {
+			excess_ref = object->excess_ref;
+			/* no need for update_refs() if object already gray */
+		} else {
+			excess_ref = 0;
+			update_refs(object);
+		}
 		spin_unlock(&object->lock);
+
+		if (excess_ref) {
+			object = lookup_object(excess_ref, 0);
+			if (!object)
+				continue;
+			if (object == scanned)
+				/* circular reference, ignore */
+				continue;
+			spin_lock_nested(&object->lock, SINGLE_DEPTH_NESTING);
+			update_refs(object);
+			spin_unlock(&object->lock);
+		}
 	}
 	read_unlock_irqrestore(&kmemleak_lock, flags);
 }
@@ -1987,6 +2068,10 @@ void __init kmemleak_init(void)
 		case KMEMLEAK_NO_SCAN:
 			kmemleak_no_scan(log->ptr);
 			break;
+		case KMEMLEAK_SET_EXCESS_REF:
+			object_set_excess_ref((unsigned long)log->ptr,
+					      log->excess_ref);
+			break;
 		default:
 			kmemleak_warn("Unknown early log operation: %d\n",
 				      log->op_type);
diff --git a/mm/vmalloc.c b/mm/vmalloc.c
index 34a1c3e46ed7..b805cc5ecca0 100644
--- a/mm/vmalloc.c
+++ b/mm/vmalloc.c
@@ -1759,12 +1759,7 @@ void *__vmalloc_node_range(unsigned long size, unsigned long align,
 	 */
 	clear_vm_uninitialized_flag(area);
 
-	/*
-	 * A ref_count = 2 is needed because vm_struct allocated in
-	 * __get_vm_area_node() contains a reference to the virtual address of
-	 * the vmalloc'ed block.
-	 */
-	kmemleak_alloc(addr, real_size, 2, gfp_mask);
+	kmemleak_vmalloc(area, size, gfp_mask);
 
 	return addr;