Re: [PATCH v3 2/5] mm: kmem: add direct objcg pointer to task_struct

[Roman Gushchin]

On Wed, Oct 18, 2023 at 11:52:27AM +0200, Vlastimil Babka wrote:
> On 10/17/23 00:18, Roman Gushchin wrote:
> > To charge a freshly allocated kernel object to a memory cgroup, the
> > kernel needs to obtain an objcg pointer. Currently it does it
> > indirectly by obtaining the memcg pointer first and then calling to
> > __get_obj_cgroup_from_memcg().
> > 
> > Usually tasks spend their entire life belonging to the same object
> > cgroup. So it makes sense to save the objcg pointer on task_struct
> > directly, so it can be obtained faster. It requires some work on fork,
> > exit and cgroup migrate paths, but these paths are way colder.
> > 
> > To avoid any costly synchronization the following rules are applied:
> > 1) A task sets it's objcg pointer itself.
> > 
> > 2) If a task is being migrated to another cgroup, the least
> >    significant bit of the objcg pointer is set atomically.
> > 
> > 3) On the allocation path the objcg pointer is obtained locklessly
> >    using the READ_ONCE() macro and the least significant bit is
> >    checked. If it's set, the following procedure is used to update
> >    it locklessly:
> >        - task->objcg is zeroed using cmpxcg
> >        - new objcg pointer is obtained
> >        - task->objcg is updated using try_cmpxchg
> >        - operation is repeated if try_cmpxcg fails
> >    It guarantees that no updates will be lost if task migration
> >    is racing against objcg pointer update. It also allows to keep
> >    both read and write paths fully lockless.
> > 
> > Because the task is keeping a reference to the objcg, it can't go away
> > while the task is alive.
> > 
> > This commit doesn't change the way the remote memcg charging works.
> > 
> > Signed-off-by: Roman Gushchin (Cruise) <roman.gushchin@xxxxxxxxx>
> > Tested-by: Naresh Kamboju <naresh.kamboju@xxxxxxxxxx>
> > Acked-by: Johannes Weiner <hannes@xxxxxxxxxxx>
> > ---
> >  include/linux/sched.h |   4 ++
> >  mm/memcontrol.c       | 130 +++++++++++++++++++++++++++++++++++++++---
> >  2 files changed, 125 insertions(+), 9 deletions(-)
> > 
> > diff --git a/mm/memcontrol.c b/mm/memcontrol.c
> > index 16ac2a5838fb..0605e45bd4a2 100644
> > --- a/mm/memcontrol.c
> > +++ b/mm/memcontrol.c
> 
> So IIUC here we increase objcg refcount.
> 
> > +				break;
> > +			objcg = NULL;
> > +		}
> > +		rcu_read_unlock();
> > +
> > +		/*
> > +		 * Try set up a new objcg pointer atomically. If it
> > +		 * fails, it means the update flag was set concurrently, so
> > +		 * the whole procedure should be repeated.
> > +		 */
> > +	} while (!try_cmpxchg(&current->objcg, &old, objcg));
> 
> And if this fails we throw objcg away and try again, but we should do
> obj_cgroup_put(objcg) first, as otherwise it would cause a leak?

Great catch! Thanks!

> 
> > +
> > +	return objcg;
> > +}
> > +
> >  __always_inline struct obj_cgroup *get_obj_cgroup_from_current(void)
> >  {
> >  	struct mem_cgroup *memcg;
> > @@ -3008,19 +3054,26 @@ __always_inline struct obj_cgroup *get_obj_cgroup_from_current(void)
> >  
> >  	if (in_task()) {
> >  		memcg = current->active_memcg;
> > +		if (unlikely(memcg))
> > +			goto from_memcg;
> >  
> > -		/* Memcg to charge can't be determined. */
> > -		if (likely(!memcg) && (!current->mm || (current->flags & PF_KTHREAD)))
> 
> The checks for current->mm and PF_KTHREAD seem to be gone completely after
> the patch, was that intended and why?

There is no need for those anymore because it's as cheap or cheaper
to check task->objcg for being NULL. Those were primarily used to rule out
kernel threads allocations early.

I gonna fix the objcg ref leak, add the comment you asked above and post v4
of this particular patch.

Thank you for reviewing the series!