cgroup LKM example

cgroup_notifier_counter.c

/*
 * cgroup-notifier-counter.c
 *
 * Author : peo3 
 *
 * Based on cgroup_freezer.c
 *
 * Copyright IBM Corporation, 2007
 *
 * Author : Cedric Le Goater <[email protected]>
 *
 * This program is free software; you can redistribute it and/or modify it
 * under the terms of version 2.1 of the GNU Lesser General Public License
 * as published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it would be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */

#include <linux/module.h>
#include <linux/slab.h>
#include <linux/cgroup.h>
#include <linux/notifier.h>
#include <linux/preempt.h>
#include <linux/kdebug.h>

extern struct cgroup_subsys notifier_subsys;
#define notifier_subsys_id notifier_subsys.subsys_id

enum notifier_types {
	NOTIFIER_PREEMPT = 0,
	NOTIFIER_DIE,
	NOTIFIER_TYPES,
};

struct notifier {
	struct cgroup_subsys_state css;
	u64 counters[NOTIFIER_TYPES];
	spinlock_t lock; /* protects _writes_ to state */
};

static inline struct notifier *cgroup_notifier(
		struct cgroup *cgroup)
{
	return container_of(
		cgroup_subsys_state(cgroup, notifier_subsys_id),
		struct notifier, css);
}

static inline struct notifier *task_notifier(struct task_struct *task)
{
	return container_of(task_subsys_state(task, notifier_subsys_id),
			    struct notifier, css);
}

static struct cgroup_subsys_state *notifier_create(struct cgroup_subsys *ss, struct cgroup *cgroup)
{
	struct notifier *notifier;
	int i;

	notifier = kzalloc(sizeof(struct notifier), GFP_KERNEL);
	if (!notifier)
		return ERR_PTR(-ENOMEM);

	spin_lock_init(&notifier->lock);
	for (i=0; i < NOTIFIER_TYPES; i++)
		notifier->counters[i] = 0;
	return &notifier->css;
}

static void notifier_destroy(struct cgroup_subsys *ss, struct cgroup *cgroup)
{
	struct notifier *notifier = cgroup_notifier(cgroup);

	kfree(notifier);
}

static u64 notifier_read_u64(struct cgroup *cgroup, struct cftype *cft)
{
	u64 ret;
	struct notifier *notifier = cgroup_notifier(cgroup);
	enum notifier_types type = cft->private;

	spin_lock(&notifier->lock);
	ret = notifier->counters[type];
	spin_unlock(&notifier->lock);

	return ret;
}

static struct cftype files[] = {
	{
		.name = "preempt",
		.private = NOTIFIER_PREEMPT,
		.read_u64 = notifier_read_u64,
	},
	{
		.name = "die",
		.private = NOTIFIER_DIE,
		.read_u64 = notifier_read_u64,
	},
};

static int notifier_populate(struct cgroup_subsys *ss, struct cgroup *cgroup)
{
	if (!cgroup->parent)
		return 0;
	return cgroup_add_files(cgroup, ss, files, ARRAY_SIZE(files));
}

struct cgroup_subsys notifier_subsys = {
	.name		= "notifier_counter",
	.create		= notifier_create,
	.destroy	= notifier_destroy,
	.populate	= notifier_populate,
	.use_id         = 1,
	.module         = THIS_MODULE,
};

static __read_mostly struct preempt_ops notifier_preempt_ops;
static struct preempt_notifier preempt_notifier;

static int
notifier_call_die(struct notifier_block *self, unsigned long cmd, void *ptr)
{
	struct notifier *notifier = task_notifier(get_current());
	struct cgroup *cgroup = task_cgroup(get_current(), notifier_subsys_id);

	if (!cgroup->parent)
		return 0;

	spin_lock(&notifier->lock);
	notifier->counters[NOTIFIER_DIE] += 1;
	spin_unlock(&notifier->lock);

	return 0;
}

static struct notifier_block notifier_block_die = {
	.notifier_call  = notifier_call_die,
};

static void notifier_sched_in(struct preempt_notifier *pn, int cpu)
{
}

static void notifier_sched_out(struct preempt_notifier *pn, struct task_struct *next)
{
	struct notifier *notifier = task_notifier(next);
	struct cgroup *cgroup = task_cgroup(next, notifier_subsys_id);

	if (!cgroup->parent)
		return;

	spin_lock(&notifier->lock);
	notifier->counters[NOTIFIER_PREEMPT] += 1;
	spin_unlock(&notifier->lock);
}

static int __init init_cgroup_notifier(void)
{
	int ret = cgroup_load_subsys(&notifier_subsys);
	if (ret)
		return ret;
	notifier_preempt_ops.sched_in = notifier_sched_in;
	notifier_preempt_ops.sched_out = notifier_sched_out;
	preempt_notifier_init(&preempt_notifier, &notifier_preempt_ops);
	preempt_notifier_register(&preempt_notifier);
	register_die_notifier(&notifier_block_die);

	return ret;
}

static void __exit exit_cgroup_notifier(void)
{
	preempt_notifier_unregister(&preempt_notifier);
	cgroup_unload_subsys(&notifier_subsys);
	unregister_die_notifier(&notifier_block_die);
}

module_init(init_cgroup_notifier);
module_exit(exit_cgroup_notifier);
MODULE_LICENSE("GPL");

Makefile

obj-m = cgroup-notifier-counter.o

KVERSION = $(shell uname -r)

all:
	make -C /lib/modules/$(KVERSION)/build M=$(PWD) modules
clean:
	make -C /lib/modules/$(KVERSION)/build M=$(PWD) clean