atomic_py.c

#include <Python.h>
#include "structmember.h"


#define atomic_int_add(ptr, val) __sync_add_and_fetch(ptr, val)

typedef struct {
    PyObject_HEAD
    int number;
} AtomicInt;

static void
AtomicInt_dealloc(AtomicInt* self)
{
    self->ob_type->tp_free((PyObject*)self);
}


static PyObject*
AtomicInt_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
{
    AtomicInt *self;

    self = (AtomicInt *)type->tp_alloc(type, 0);
    if (self != NULL) {
        self->number = 0;
    }

    return (PyObject *)self;
}

static int
AtomicInt_init(AtomicInt *self, PyObject *args)
{
    if (! PyArg_ParseTuple(args, "i", &self->number))
        return -1;
    return 0;
}

static PyMemberDef AtomicInt_members[] = {
    {"number", T_INT, offsetof(AtomicInt, number), 0, "atomic number"},
    {NULL}  /* Sentinel */
};


PyObject*
AtomicInt_increment(AtomicInt *self)
{
    Py_INCREF(self);
    atomic_int_add(&self->number, 1);
    Py_DECREF(self);
    return Py_BuildValue("");
}

static PyMethodDef AtomicInt_methods[] = {
    {"increment", (PyCFunction)AtomicInt_increment, METH_NOARGS, "Atomically increment an integer"},
    {NULL}
};

static PyTypeObject AtomicIntType = {
    PyObject_HEAD_INIT(NULL)
    0,                         /*ob_size*/
    "atomic.AtomicInt",             /*tp_name*/
    sizeof(AtomicInt),             /*tp_basicsize*/
    0,                         /*tp_itemsize*/
    (destructor)AtomicInt_dealloc, /*tp_dealloc*/
    0,                         /*tp_print*/
    0,                         /*tp_getattr*/
    0,                         /*tp_setattr*/
    0,                         /*tp_compare*/
    0,                         /*tp_repr*/
    0,                         /*tp_as_number*/
    0,                         /*tp_as_sequence*/
    0,                         /*tp_as_mapping*/
    0,                         /*tp_hash */
    0,                         /*tp_call*/
    0,                         /*tp_str*/
    0,                         /*tp_getattro*/
    0,                         /*tp_setattro*/
    0,                         /*tp_as_buffer*/
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /*tp_flags*/
    "Atomic Int objects",           /* tp_doc */
    0,		               /* tp_traverse */
    0,		               /* tp_clear */
    0,		               /* tp_richcompare */
    0,		               /* tp_weaklistoffset */
    0,		               /* tp_iter */
    0,		               /* tp_iternext */
    AtomicInt_methods,             /* tp_methods */
    AtomicInt_members,             /* tp_members */
    0,                         /* tp_getset */
    0,                         /* tp_base */
    0,                         /* tp_dict */
    0,                         /* tp_descr_get */
    0,                         /* tp_descr_set */
    0,                         /* tp_dictoffset */
    (initproc)AtomicInt_init,      /* tp_init */
    0,                         /* tp_alloc */
    AtomicInt_new,                 /* tp_new */
};


static PyMethodDef module_methods[] = {
    {NULL}
};


PyMODINIT_FUNC
initatomic(void)
{
    PyObject* m = NULL;

    if (PyType_Ready(&AtomicIntType) < 0)
        return;
    m = Py_InitModule3("atomic", module_methods, "why are we here");
    if (m == NULL)
	return;
    
    Py_INCREF(&AtomicIntType);
    PyModule_AddObject(m, "AtomicInt", (PyObject *)&AtomicIntType);
}

setup.py

from setuptools import setup, Extension

atomic = Extension('atomic',
                   sources=['atomic_py.c'])

setup(name='AtomicPy',
      version='dev',
      ext_modules=[atomic]
      )

test.py

"""
Run test.py with 1 for atomic, 0 for non-AtomicWorker
"""
import sys
import threading

import atomic


atom = None


class Worker(threading.Thread):
    def run(self):
        global atom
        for _ in range(500):
            atom += 1
            # with lock:
            #     print foo


class AtomicWorker(threading.Thread):
    def run(self):
        global atom
        for _ in range(500):
            atom.increment()
            # foo = atomic.int_add(foo, 1)
            # with lock:
            #     print foo


def main(use_atomic):
    workers = {}
    global atom
    if use_atomic:
        atom = atomic.AtomicInt(0)
    else:
        atom = 0
    for x in range(10):
        if use_atomic:
            thread = AtomicWorker()
        else:
            thread = Worker()
        thread.start()
        workers[x] = thread
    for _, worker in workers.items():
        worker.join()
    global foo
    # assert foo == 500, "Derp"
    if use_atomic:
        print atom.number
    else:
        print atom

if __name__ == '__main__':
    sys.setcheckinterval(0)
    main(eval(sys.argv[1]))