diff mbox series

[v3] trace-cruncher: Add API to set tracing CPU affinity

Message ID 20220117201835.0991759f@gandalf.local.home (mailing list archive)
State New, archived
Headers show
Series [v3] trace-cruncher: Add API to set tracing CPU affinity | expand

Commit Message

Steven Rostedt Jan. 18, 2022, 1:18 a.m. UTC
From: "Steven Rostedt (VMware)" <rostedt@goodmis.org>

Add a set_affinity API that can let the user set what CPUs to enable
tracing on.

Example:

  import tracecruncher.ftracepy as ft

  ft.set_affinity(cpus=["0-1","3"]);

The above will set the top level instance affinity to 0,1,3 (or 0xb)

To reset back to online CPUs:

  ft.reset_affinity()

Where it will reset the top level instance back to all CPUS.

Signed-off-by: Steven Rostedt (VMware) <rostedt@goodmis.org>
---
Changes since v2: https://lore.kernel.org/all/20220117201231.414a4799@gandalf.local.home/

 - Fix (again) the description of set_affinity.

 examples/reset-affinity.py |  5 ++
 examples/test-affinity.py  |  5 ++
 src/ftracepy-utils.c       | 93 ++++++++++++++++++++++++++++++++++++++
 src/ftracepy-utils.h       |  6 +++
 src/ftracepy.c             | 10 ++++
 5 files changed, 119 insertions(+)
 create mode 100755 examples/reset-affinity.py
 create mode 100755 examples/test-affinity.py

Comments

Yordan Karadzhov Jan. 18, 2022, 10:08 a.m. UTC | #1
On 18.01.22 г. 3:18 ч., Steven Rostedt wrote:
> From: "Steven Rostedt (VMware)" <rostedt@goodmis.org>
> 
> Add a set_affinity API that can let the user set what CPUs to enable
> tracing on.
> 
> Example:
> 
>    import tracecruncher.ftracepy as ft
> 
>    ft.set_affinity(cpus=["0-1","3"]);
> 
> The above will set the top level instance affinity to 0,1,3 (or 0xb)
> 
> To reset back to online CPUs:
> 
>    ft.reset_affinity()
> 
> Where it will reset the top level instance back to all CPUS.
> 
> Signed-off-by: Steven Rostedt (VMware) <rostedt@goodmis.org>
> ---
> Changes since v2: https://lore.kernel.org/all/20220117201231.414a4799@gandalf.local.home/
> 
>   - Fix (again) the description of set_affinity.
> 
>   examples/reset-affinity.py |  5 ++
>   examples/test-affinity.py  |  5 ++
>   src/ftracepy-utils.c       | 93 ++++++++++++++++++++++++++++++++++++++
>   src/ftracepy-utils.h       |  6 +++
>   src/ftracepy.c             | 10 ++++
>   5 files changed, 119 insertions(+)
>   create mode 100755 examples/reset-affinity.py
>   create mode 100755 examples/test-affinity.py
> 
> diff --git a/examples/reset-affinity.py b/examples/reset-affinity.py
> new file mode 100755
> index 0000000..2bca32e
> --- /dev/null
> +++ b/examples/reset-affinity.py
> @@ -0,0 +1,5 @@
> +#!/usr/bin/env python3
> +
> +import tracecruncher.ftracepy as ft
> +
> +ft.reset_affinity();
> diff --git a/examples/test-affinity.py b/examples/test-affinity.py
> new file mode 100755
> index 0000000..1619931
> --- /dev/null
> +++ b/examples/test-affinity.py
> @@ -0,0 +1,5 @@
> +#!/usr/bin/env python3
> +
> +import tracecruncher.ftracepy as ft
> +
> +ft.set_affinity(cpus=["0-1","3"]);
> diff --git a/src/ftracepy-utils.c b/src/ftracepy-utils.c
> index a735d88..e42f69c 100644
> --- a/src/ftracepy-utils.c
> +++ b/src/ftracepy-utils.c
> @@ -2394,6 +2394,99 @@ PyObject *PyFtrace_hist(PyObject *self, PyObject *args,
>   	return NULL;
>   }
>   
> +PyObject *PyFtrace_reset_affinity(PyObject *self, PyObject *args,
> +				  PyObject *kwargs)
> +{
> +	struct tracefs_instance *instance;
> +	static char *kwlist[] = {"instance", NULL};
> +	PyObject *py_inst = NULL;
> +	int ret;
> +
> +	if (!PyArg_ParseTupleAndKeywords(args,
> +					 kwargs,
> +					 "|O",
> +					 kwlist,
> +					 &py_inst)) {
> +		return NULL;
> +	}
> +
> +	if (!get_optional_instance(py_inst, &instance))
> +		goto err;

No need to have 'goto' in this function. Just return NULL (see also my comment at the bottom).


> +
> +	/* NULL for CPUs will set all available CPUS */
> +	ret = tracefs_instance_set_affinity(instance, NULL);
> +	if (ret < 0) {
> +		PyErr_SetString(TRACECRUNCHER_ERROR,
> +				"Failed to reset tracing affinity.");
> +		goto err;
> +	}

Can we make the call of tracefs_instance_set_affinity(), the error check and the error message a static helper?
I will explain why bellow in the patch.

> +
> +	Py_RETURN_NONE;
> +err:
> +	return NULL;
> +}
> +
> +PyObject *PyFtrace_set_affinity(PyObject *self, PyObject *args,
> +						PyObject *kwargs)
> +{
> +	struct tracefs_instance *instance;
> +	static char *kwlist[] = {"cpus", "instance", NULL};
> +	PyObject *py_cpus;
> +	PyObject *py_inst = NULL;
> +	const char *cpu_str;
> +	int ret;
> +
> +	if (!init_print_seq())
> +		return false;

We have to return NULL here.

> +
> +	if (!PyArg_ParseTupleAndKeywords(args,
> +					 kwargs,
> +					 "O|O",
> +					 kwlist,
> +					 &py_cpus,
> +					 &py_inst)) {
> +		return NULL;
> +	}
> +
> +	if (PyUnicode_Check(py_cpus)) {
> +		cpu_str = (const char *)PyUnicode_DATA(py_cpus);

And here we can have

		if ((is_all(cpu_str) {
			Call the static helper for 'reset' here and return.
		}

> +		if (trace_seq_puts(&seq, cpu_str) < > +			goto err_seq;
> +	} else if (PyList_CheckExact(py_cpus)) {
> +		int i, n = PyList_Size(py_cpus);
> +
> +		for (i = 0; i < n; ++i) {
> +			cpu_str = str_from_list(py_cpus, i);
> +			if (i)
> +				trace_seq_putc(&seq, ',');
> +			if (trace_seq_puts(&seq, cpu_str) < 0)
> +				goto err_seq;
> +		}
> +	} else {
> +		PyErr_SetString(TRACECRUNCHER_ERROR,
> +				"Invalid \"cpus\" type.");

It may be better to use TfsError_fmt() here.
This helper will print also the error log of libtracefs.
The same comment applies for all error messages below.

> +		goto err;
  		return NULL;
> +	}
> +
> +	if (!get_optional_instance(py_inst, &instance))
> +		goto err;
		return NULL;
> +
> +	trace_seq_terminate(&seq);
> +	ret = tracefs_instance_set_affinity(instance, seq.buffer);
> +	if (ret < 0) {
> +		PyErr_SetString(TRACECRUNCHER_ERROR,
> +				"Invalid \"cpus\" argument.");
> +		goto err;
  		return NULL;
> +	}
> +
> +	Py_RETURN_NONE;
> +err_seq:
> +	PyErr_SetString(TRACECRUNCHER_ERROR,
> +			"Internal processing string error.");
> +err:
> +	return NULL;

I guess you use 'goto err;' everywhere, because you want to stress that returning NULL is essentially exiting with 
error. However, so far we used just 'return NULL' everywhere and I would prefer to keep the code consistent.
Maybe in a separate patch we can replace all ''return NULL;' lines with a macro that can have some self-explanatory name?

Thanks!
Yordan




> +}
> +
>   PyObject *PyFtrace_set_ftrace_loglevel(PyObject *self, PyObject *args,
>   						       PyObject *kwargs)
>   {
> diff --git a/src/ftracepy-utils.h b/src/ftracepy-utils.h
> index d09c8bf..4af557f 100644
> --- a/src/ftracepy-utils.h
> +++ b/src/ftracepy-utils.h
> @@ -208,6 +208,12 @@ PyObject *PyFtrace_register_kretprobe(PyObject *self, PyObject *args,
>   PyObject *PyFtrace_hist(PyObject *self, PyObject *args,
>   					PyObject *kwargs);
>   
> +PyObject *PyFtrace_set_affinity(PyObject *self, PyObject *args,
> +						PyObject *kwargs);
> +
> +PyObject *PyFtrace_reset_affinity(PyObject *self, PyObject *args,
> +						  PyObject *kwargs);
> +
>   PyObject *PyFtrace_set_ftrace_loglevel(PyObject *self, PyObject *args,
>   						       PyObject *kwargs);
>   
> diff --git a/src/ftracepy.c b/src/ftracepy.c
> index b270b71..168ad77 100644
> --- a/src/ftracepy.c
> +++ b/src/ftracepy.c
> @@ -372,11 +372,21 @@ static PyMethodDef ftracepy_methods[] = {
>   	 METH_VARARGS | METH_KEYWORDS,
>   	 "Define a kretprobe."
>   	},
> +	{"reset_affinity",
> +	 (PyCFunction) PyFtrace_reset_affinity,
> +	 METH_VARARGS | METH_KEYWORDS,
> +	 "Set an instance tracing affinity to all CPUs"
> +	},
>   	{"hist",
>   	 (PyCFunction) PyFtrace_hist,
>   	 METH_VARARGS | METH_KEYWORDS,
>   	 "Define a histogram."
>   	},
> +	{"set_affinity",
> +	 (PyCFunction) PyFtrace_set_affinity,
> +	 METH_VARARGS | METH_KEYWORDS,
> +	 "Set the tracing affinity of an instance."
> +	},
>   	{"set_ftrace_loglevel",
>   	 (PyCFunction) PyFtrace_set_ftrace_loglevel,
>   	 METH_VARARGS | METH_KEYWORDS,
>
Steven Rostedt Jan. 18, 2022, 10:40 p.m. UTC | #2
On Tue, 18 Jan 2022 12:08:44 +0200
Yordan Karadzhov <y.karadz@gmail.com> wrote:

> On 18.01.22 г. 3:18 ч., Steven Rostedt wrote:

> > diff --git a/examples/reset-affinity.py b/examples/reset-affinity.py
> > new file mode 100755
> > index 0000000..2bca32e
> > --- /dev/null
> > +++ b/examples/reset-affinity.py
> > @@ -0,0 +1,5 @@
> > +#!/usr/bin/env python3
> > +
> > +import tracecruncher.ftracepy as ft
> > +
> > +ft.reset_affinity();
> > diff --git a/examples/test-affinity.py b/examples/test-affinity.py
> > new file mode 100755
> > index 0000000..1619931
> > --- /dev/null
> > +++ b/examples/test-affinity.py
> > @@ -0,0 +1,5 @@
> > +#!/usr/bin/env python3
> > +
> > +import tracecruncher.ftracepy as ft
> > +
> > +ft.set_affinity(cpus=["0-1","3"]);
> > diff --git a/src/ftracepy-utils.c b/src/ftracepy-utils.c
> > index a735d88..e42f69c 100644
> > --- a/src/ftracepy-utils.c
> > +++ b/src/ftracepy-utils.c
> > @@ -2394,6 +2394,99 @@ PyObject *PyFtrace_hist(PyObject *self, PyObject *args,
> >   	return NULL;
> >   }
> >   
> > +PyObject *PyFtrace_reset_affinity(PyObject *self, PyObject *args,
> > +				  PyObject *kwargs)
> > +{
> > +	struct tracefs_instance *instance;
> > +	static char *kwlist[] = {"instance", NULL};
> > +	PyObject *py_inst = NULL;
> > +	int ret;
> > +
> > +	if (!PyArg_ParseTupleAndKeywords(args,
> > +					 kwargs,
> > +					 "|O",
> > +					 kwlist,
> > +					 &py_inst)) {
> > +		return NULL;
> > +	}
> > +
> > +	if (!get_optional_instance(py_inst, &instance))
> > +		goto err;  
> 
> No need to have 'goto' in this function. Just return NULL (see also my comment at the bottom).
> 
> 
> > +
> > +	/* NULL for CPUs will set all available CPUS */
> > +	ret = tracefs_instance_set_affinity(instance, NULL);
> > +	if (ret < 0) {
> > +		PyErr_SetString(TRACECRUNCHER_ERROR,
> > +				"Failed to reset tracing affinity.");
> > +		goto err;
> > +	}  
> 
> Can we make the call of tracefs_instance_set_affinity(), the error check and the error message a static helper?
> I will explain why bellow in the patch.
> 
> > +
> > +	Py_RETURN_NONE;
> > +err:
> > +	return NULL;
> > +}
> > +
> > +PyObject *PyFtrace_set_affinity(PyObject *self, PyObject *args,
> > +						PyObject *kwargs)
> > +{
> > +	struct tracefs_instance *instance;
> > +	static char *kwlist[] = {"cpus", "instance", NULL};
> > +	PyObject *py_cpus;
> > +	PyObject *py_inst = NULL;
> > +	const char *cpu_str;
> > +	int ret;
> > +
> > +	if (!init_print_seq())
> > +		return false;  
> 
> We have to return NULL here.

Oops, silly cut and paste issue.

> 
> > +
> > +	if (!PyArg_ParseTupleAndKeywords(args,
> > +					 kwargs,
> > +					 "O|O",
> > +					 kwlist,
> > +					 &py_cpus,
> > +					 &py_inst)) {
> > +		return NULL;
> > +	}
> > +
> > +	if (PyUnicode_Check(py_cpus)) {
> > +		cpu_str = (const char *)PyUnicode_DATA(py_cpus);  
> 
> And here we can have
> 
> 		if ((is_all(cpu_str) {
> 			Call the static helper for 'reset' here and return.

Do we need a static helper? It is simply:

	tracefs_instance_set_affinity(instance, NULL);

and that will set all of them. (NULL for CPUs means all CPUs).

> 		}
> 
> > +		if (trace_seq_puts(&seq, cpu_str) < > +			goto err_seq;
> > +	} else if (PyList_CheckExact(py_cpus)) {
> > +		int i, n = PyList_Size(py_cpus);
> > +
> > +		for (i = 0; i < n; ++i) {
> > +			cpu_str = str_from_list(py_cpus, i);
> > +			if (i)
> > +				trace_seq_putc(&seq, ',');
> > +			if (trace_seq_puts(&seq, cpu_str) < 0)
> > +				goto err_seq;
> > +		}
> > +	} else {
> > +		PyErr_SetString(TRACECRUNCHER_ERROR,
> > +				"Invalid \"cpus\" type.");  
> 
> It may be better to use TfsError_fmt() here.
> This helper will print also the error log of libtracefs.
> The same comment applies for all error messages below.

Do we want that as the error log will not be modified by this operation.

> 
> > +		goto err;  
>   		return NULL;
> > +	}
> > +
> > +	if (!get_optional_instance(py_inst, &instance))
> > +		goto err;  
> 		return NULL;
> > +
> > +	trace_seq_terminate(&seq);
> > +	ret = tracefs_instance_set_affinity(instance, seq.buffer);
> > +	if (ret < 0) {
> > +		PyErr_SetString(TRACECRUNCHER_ERROR,
> > +				"Invalid \"cpus\" argument.");
> > +		goto err;  
>   		return NULL;
> > +	}
> > +
> > +	Py_RETURN_NONE;
> > +err_seq:
> > +	PyErr_SetString(TRACECRUNCHER_ERROR,
> > +			"Internal processing string error.");
> > +err:
> > +	return NULL;  
> 
> I guess you use 'goto err;' everywhere, because you want to stress that returning NULL is essentially exiting with 
> error. However, so far we used just 'return NULL' everywhere and I would prefer to keep the code consistent.
> Maybe in a separate patch we can replace all ''return NULL;' lines with a macro that can have some self-explanatory name?

Yeah, I was trying to match the existing code, as I'm new to this. ;-)

-- Steve
Steven Rostedt Jan. 18, 2022, 10:41 p.m. UTC | #3
On Tue, 18 Jan 2022 17:40:19 -0500
Steven Rostedt <rostedt@goodmis.org> wrote:

> > I guess you use 'goto err;' everywhere, because you want to stress that returning NULL is essentially exiting with 
> > error. However, so far we used just 'return NULL' everywhere and I would prefer to keep the code consistent.
> > Maybe in a separate patch we can replace all ''return NULL;' lines with a macro that can have some self-explanatory name?  
> 
> Yeah, I was trying to match the existing code, as I'm new to this. ;-)

Oh, if you haven't noticed. I created a tracefs_instance_get_affinity() API
as well ;-)

  https://patchwork.kernel.org/project/linux-trace-devel/patch/20220118030005.1603821-2-rostedt@goodmis.org/

-- Steve
Yordan Karadzhov Jan. 19, 2022, 10:24 a.m. UTC | #4
On 19.01.22 г. 0:40 ч., Steven Rostedt wrote:
> On Tue, 18 Jan 2022 12:08:44 +0200
> Yordan Karadzhov <y.karadz@gmail.com> wrote:
> 
>> On 18.01.22 г. 3:18 ч., Steven Rostedt wrote:
> 
>>> diff --git a/examples/reset-affinity.py b/examples/reset-affinity.py
>>> new file mode 100755
>>> index 0000000..2bca32e
>>> --- /dev/null
>>> +++ b/examples/reset-affinity.py
>>> @@ -0,0 +1,5 @@
>>> +#!/usr/bin/env python3
>>> +
>>> +import tracecruncher.ftracepy as ft
>>> +
>>> +ft.reset_affinity();
>>> diff --git a/examples/test-affinity.py b/examples/test-affinity.py
>>> new file mode 100755
>>> index 0000000..1619931
>>> --- /dev/null
>>> +++ b/examples/test-affinity.py
>>> @@ -0,0 +1,5 @@
>>> +#!/usr/bin/env python3
>>> +
>>> +import tracecruncher.ftracepy as ft
>>> +
>>> +ft.set_affinity(cpus=["0-1","3"]);
>>> diff --git a/src/ftracepy-utils.c b/src/ftracepy-utils.c
>>> index a735d88..e42f69c 100644
>>> --- a/src/ftracepy-utils.c
>>> +++ b/src/ftracepy-utils.c
>>> @@ -2394,6 +2394,99 @@ PyObject *PyFtrace_hist(PyObject *self, PyObject *args,
>>>    	return NULL;
>>>    }
>>>    
>>> +PyObject *PyFtrace_reset_affinity(PyObject *self, PyObject *args,
>>> +				  PyObject *kwargs)
>>> +{
>>> +	struct tracefs_instance *instance;
>>> +	static char *kwlist[] = {"instance", NULL};
>>> +	PyObject *py_inst = NULL;
>>> +	int ret;
>>> +
>>> +	if (!PyArg_ParseTupleAndKeywords(args,
>>> +					 kwargs,
>>> +					 "|O",
>>> +					 kwlist,
>>> +					 &py_inst)) {
>>> +		return NULL;
>>> +	}
>>> +
>>> +	if (!get_optional_instance(py_inst, &instance))
>>> +		goto err;
>>
>> No need to have 'goto' in this function. Just return NULL (see also my comment at the bottom).
>>
>>
>>> +
>>> +	/* NULL for CPUs will set all available CPUS */
>>> +	ret = tracefs_instance_set_affinity(instance, NULL);
>>> +	if (ret < 0) {
>>> +		PyErr_SetString(TRACECRUNCHER_ERROR,
>>> +				"Failed to reset tracing affinity.");
>>> +		goto err;
>>> +	}
>>
>> Can we make the call of tracefs_instance_set_affinity(), the error check and the error message a static helper?
>> I will explain why bellow in the patch.
>>
>>> +
>>> +	Py_RETURN_NONE;
>>> +err:
>>> +	return NULL;
>>> +}
>>> +
>>> +PyObject *PyFtrace_set_affinity(PyObject *self, PyObject *args,
>>> +						PyObject *kwargs)
>>> +{
>>> +	struct tracefs_instance *instance;
>>> +	static char *kwlist[] = {"cpus", "instance", NULL};
>>> +	PyObject *py_cpus;
>>> +	PyObject *py_inst = NULL;
>>> +	const char *cpu_str;
>>> +	int ret;
>>> +
>>> +	if (!init_print_seq())
>>> +		return false;
>>
>> We have to return NULL here.
> 
> Oops, silly cut and paste issue.
> 
>>
>>> +
>>> +	if (!PyArg_ParseTupleAndKeywords(args,
>>> +					 kwargs,
>>> +					 "O|O",
>>> +					 kwlist,
>>> +					 &py_cpus,
>>> +					 &py_inst)) {
>>> +		return NULL;
>>> +	}
>>> +
>>> +	if (PyUnicode_Check(py_cpus)) {
>>> +		cpu_str = (const char *)PyUnicode_DATA(py_cpus);
>>
>> And here we can have
>>
>> 		if ((is_all(cpu_str) {
>> 			Call the static helper for 'reset' here and return.
> 
> Do we need a static helper? It is simply:
> 
> 	tracefs_instance_set_affinity(instance, NULL);
> 
> and that will set all of them. (NULL for CPUs means all CPUs).

Yes, but we still need to check the return value and print an error message in the case of a failure.

> 
>> 		}
>>
>>> +		if (trace_seq_puts(&seq, cpu_str) < > +			goto err_seq;
>>> +	} else if (PyList_CheckExact(py_cpus)) {
>>> +		int i, n = PyList_Size(py_cpus);
>>> +
>>> +		for (i = 0; i < n; ++i) {
>>> +			cpu_str = str_from_list(py_cpus, i);
>>> +			if (i)
>>> +				trace_seq_putc(&seq, ',');
>>> +			if (trace_seq_puts(&seq, cpu_str) < 0)
>>> +				goto err_seq;
>>> +		}
>>> +	} else {
>>> +		PyErr_SetString(TRACECRUNCHER_ERROR,
>>> +				"Invalid \"cpus\" type.");
>>
>> It may be better to use TfsError_fmt() here.
>> This helper will print also the error log of libtracefs.
>> The same comment applies for all error messages below.
> 
> Do we want that as the error log will not be modified by this operation.

So far I'v been trying to print the error log every time when a libtracefs API returns an error. I know that not all of 
the errors are adding messages to the log, but I was thinking that this is something that is not guarantied to be stable 
and may change in the future.

Thanks!
Yordan

> 
>>
>>> +		goto err;
>>    		return NULL;
>>> +	}
>>> +
>>> +	if (!get_optional_instance(py_inst, &instance))
>>> +		goto err;
>> 		return NULL;
>>> +
>>> +	trace_seq_terminate(&seq);
>>> +	ret = tracefs_instance_set_affinity(instance, seq.buffer);
>>> +	if (ret < 0) {
>>> +		PyErr_SetString(TRACECRUNCHER_ERROR,
>>> +				"Invalid \"cpus\" argument.");
>>> +		goto err;
>>    		return NULL;
>>> +	}
>>> +
>>> +	Py_RETURN_NONE;
>>> +err_seq:
>>> +	PyErr_SetString(TRACECRUNCHER_ERROR,
>>> +			"Internal processing string error.");
>>> +err:
>>> +	return NULL;
>>
>> I guess you use 'goto err;' everywhere, because you want to stress that returning NULL is essentially exiting with
>> error. However, so far we used just 'return NULL' everywhere and I would prefer to keep the code consistent.
>> Maybe in a separate patch we can replace all ''return NULL;' lines with a macro that can have some self-explanatory name?
> 
> Yeah, I was trying to match the existing code, as I'm new to this. ;-)
> 
> -- Steve
>
Yordan Karadzhov Jan. 19, 2022, 10:30 a.m. UTC | #5
On 19.01.22 г. 0:41 ч., Steven Rostedt wrote:
> On Tue, 18 Jan 2022 17:40:19 -0500
> Steven Rostedt <rostedt@goodmis.org> wrote:
> 
>>> I guess you use 'goto err;' everywhere, because you want to stress that returning NULL is essentially exiting with
>>> error. However, so far we used just 'return NULL' everywhere and I would prefer to keep the code consistent.
>>> Maybe in a separate patch we can replace all ''return NULL;' lines with a macro that can have some self-explanatory name?
>>
>> Yeah, I was trying to match the existing code, as I'm new to this. ;-)
> 
> Oh, if you haven't noticed. I created a tracefs_instance_get_affinity() API
> as well ;-)
> 
>    https://patchwork.kernel.org/project/linux-trace-devel/patch/20220118030005.1603821-2-rostedt@goodmis.org/

Great!
Y.

> 
> -- Steve
>
diff mbox series

Patch

diff --git a/examples/reset-affinity.py b/examples/reset-affinity.py
new file mode 100755
index 0000000..2bca32e
--- /dev/null
+++ b/examples/reset-affinity.py
@@ -0,0 +1,5 @@ 
+#!/usr/bin/env python3
+
+import tracecruncher.ftracepy as ft
+
+ft.reset_affinity();
diff --git a/examples/test-affinity.py b/examples/test-affinity.py
new file mode 100755
index 0000000..1619931
--- /dev/null
+++ b/examples/test-affinity.py
@@ -0,0 +1,5 @@ 
+#!/usr/bin/env python3
+
+import tracecruncher.ftracepy as ft
+
+ft.set_affinity(cpus=["0-1","3"]);
diff --git a/src/ftracepy-utils.c b/src/ftracepy-utils.c
index a735d88..e42f69c 100644
--- a/src/ftracepy-utils.c
+++ b/src/ftracepy-utils.c
@@ -2394,6 +2394,99 @@  PyObject *PyFtrace_hist(PyObject *self, PyObject *args,
 	return NULL;
 }
 
+PyObject *PyFtrace_reset_affinity(PyObject *self, PyObject *args,
+				  PyObject *kwargs)
+{
+	struct tracefs_instance *instance;
+	static char *kwlist[] = {"instance", NULL};
+	PyObject *py_inst = NULL;
+	int ret;
+
+	if (!PyArg_ParseTupleAndKeywords(args,
+					 kwargs,
+					 "|O",
+					 kwlist,
+					 &py_inst)) {
+		return NULL;
+	}
+
+	if (!get_optional_instance(py_inst, &instance))
+		goto err;
+
+	/* NULL for CPUs will set all available CPUS */
+	ret = tracefs_instance_set_affinity(instance, NULL);
+	if (ret < 0) {
+		PyErr_SetString(TRACECRUNCHER_ERROR,
+				"Failed to reset tracing affinity.");
+		goto err;
+	}
+
+	Py_RETURN_NONE;
+err:
+	return NULL;
+}
+
+PyObject *PyFtrace_set_affinity(PyObject *self, PyObject *args,
+						PyObject *kwargs)
+{
+	struct tracefs_instance *instance;
+	static char *kwlist[] = {"cpus", "instance", NULL};
+	PyObject *py_cpus;
+	PyObject *py_inst = NULL;
+	const char *cpu_str;
+	int ret;
+
+	if (!init_print_seq())
+		return false;
+
+	if (!PyArg_ParseTupleAndKeywords(args,
+					 kwargs,
+					 "O|O",
+					 kwlist,
+					 &py_cpus,
+					 &py_inst)) {
+		return NULL;
+	}
+
+	if (PyUnicode_Check(py_cpus)) {
+		cpu_str = (const char *)PyUnicode_DATA(py_cpus);
+		if (trace_seq_puts(&seq, cpu_str) < 0)
+			goto err_seq;
+	} else if (PyList_CheckExact(py_cpus)) {
+		int i, n = PyList_Size(py_cpus);
+
+		for (i = 0; i < n; ++i) {
+			cpu_str = str_from_list(py_cpus, i);
+			if (i)
+				trace_seq_putc(&seq, ',');
+			if (trace_seq_puts(&seq, cpu_str) < 0)
+				goto err_seq;
+		}
+	} else {
+		PyErr_SetString(TRACECRUNCHER_ERROR,
+				"Invalid \"cpus\" type.");
+		goto err;
+	}
+
+	if (!get_optional_instance(py_inst, &instance))
+		goto err;
+
+	trace_seq_terminate(&seq);
+	ret = tracefs_instance_set_affinity(instance, seq.buffer);
+	if (ret < 0) {
+		PyErr_SetString(TRACECRUNCHER_ERROR,
+				"Invalid \"cpus\" argument.");
+		goto err;
+	}
+
+	Py_RETURN_NONE;
+err_seq:
+	PyErr_SetString(TRACECRUNCHER_ERROR,
+			"Internal processing string error.");
+err:
+	return NULL;
+}
+
 PyObject *PyFtrace_set_ftrace_loglevel(PyObject *self, PyObject *args,
 						       PyObject *kwargs)
 {
diff --git a/src/ftracepy-utils.h b/src/ftracepy-utils.h
index d09c8bf..4af557f 100644
--- a/src/ftracepy-utils.h
+++ b/src/ftracepy-utils.h
@@ -208,6 +208,12 @@  PyObject *PyFtrace_register_kretprobe(PyObject *self, PyObject *args,
 PyObject *PyFtrace_hist(PyObject *self, PyObject *args,
 					PyObject *kwargs);
 
+PyObject *PyFtrace_set_affinity(PyObject *self, PyObject *args,
+						PyObject *kwargs);
+
+PyObject *PyFtrace_reset_affinity(PyObject *self, PyObject *args,
+						  PyObject *kwargs);
+
 PyObject *PyFtrace_set_ftrace_loglevel(PyObject *self, PyObject *args,
 						       PyObject *kwargs);
 
diff --git a/src/ftracepy.c b/src/ftracepy.c
index b270b71..168ad77 100644
--- a/src/ftracepy.c
+++ b/src/ftracepy.c
@@ -372,11 +372,21 @@  static PyMethodDef ftracepy_methods[] = {
 	 METH_VARARGS | METH_KEYWORDS,
 	 "Define a kretprobe."
 	},
+	{"reset_affinity",
+	 (PyCFunction) PyFtrace_reset_affinity,
+	 METH_VARARGS | METH_KEYWORDS,
+	 "Set an instance tracing affinity to all CPUs"
+	},
 	{"hist",
 	 (PyCFunction) PyFtrace_hist,
 	 METH_VARARGS | METH_KEYWORDS,
 	 "Define a histogram."
 	},
+	{"set_affinity",
+	 (PyCFunction) PyFtrace_set_affinity,
+	 METH_VARARGS | METH_KEYWORDS,
+	 "Set the tracing affinity of an instance."
+	},
 	{"set_ftrace_loglevel",
 	 (PyCFunction) PyFtrace_set_ftrace_loglevel,
 	 METH_VARARGS | METH_KEYWORDS,