#include "Python.h"

#include "interpreteridobject.h"

#define MODULE_NAME "_xxinterpchannels"

#define GLOBAL_MALLOC(TYPE) \

PyMem_RawMalloc(sizeof(TYPE))

#define GLOBAL_FREE(VAR) \

PyMem_RawFree(VAR)

static PyInterpreterState *

_get_current_interp(void)

{

// PyInterpreterState_Get() aborts if lookup fails, so don't need

// to check the result for NULL.

return PyInterpreterState_Get();

}

static PyObject *

_get_current_module(void)

{

PyObject *name = PyUnicode_FromString(MODULE_NAME);

if (name == NULL) {

return NULL;

}

PyObject *mod = PyImport_GetModule(name);

Py_DECREF(name);

if (mod == NULL) {

return NULL;

}

assert(mod != Py_None);

return mod;

}

static PyObject *

get_module_from_owned_type(PyTypeObject *cls)

{

assert(cls != NULL);

return _get_current_module();

// XXX Use the more efficient API now that we use heap types:

//return PyType_GetModule(cls);

}

static struct PyModuleDef moduledef;

static PyObject *

get_module_from_type(PyTypeObject *cls)

{

assert(cls != NULL);

return _get_current_module();

// XXX Use the more efficient API now that we use heap types:

//return PyType_GetModuleByDef(cls, &moduledef);

}

static PyObject *

add_new_exception(PyObject *mod, const char *name, PyObject *base)

{

assert(!PyObject_HasAttrString(mod, name));

PyObject *exctype = PyErr_NewException(name, base, NULL);

if (exctype == NULL) {

return NULL;

}

int res = PyModule_AddType(mod, (PyTypeObject *)exctype);

if (res < 0) {

Py_DECREF(exctype);

return NULL;

}

return exctype;

}

#define ADD_NEW_EXCEPTION(MOD, NAME, BASE) \

add_new_exception(MOD, MODULE_NAME "." Py_STRINGIFY(NAME), BASE)

static PyTypeObject *

add_new_type(PyObject *mod, PyType_Spec *spec, crossinterpdatafunc shared)

{

PyTypeObject *cls = (PyTypeObject *)PyType_FromMetaclass(

NULL, mod, spec, NULL);

if (cls == NULL) {

return NULL;

}

if (PyModule_AddType(mod, cls) < 0) {

Py_DECREF(cls);

return NULL;

}

if (shared != NULL) {

if (_PyCrossInterpreterData_RegisterClass(cls, shared)) {

Py_DECREF(cls);

return NULL;

}

}

return cls;

}

static int

_release_xid_data(_PyCrossInterpreterData *data, int ignoreexc)

{

PyObject *exc;

if (ignoreexc) {

exc = PyErr_GetRaisedException();

}

int res = _PyCrossInterpreterData_Release(data);

if (res < 0) {

/* The owning interpreter is already destroyed. */

if (ignoreexc) {

// XXX Emit a warning?

PyErr_Clear();

}

}

if (ignoreexc) {

PyErr_SetRaisedException(exc);

}

return res;

}

typedef struct {

/* heap types */

PyTypeObject *ChannelIDType;

/* exceptions */

PyObject *ChannelError;

PyObject *ChannelNotFoundError;

PyObject *ChannelClosedError;

PyObject *ChannelEmptyError;

PyObject *ChannelNotEmptyError;

} module_state;

static inline module_state *

get_module_state(PyObject *mod)

{

assert(mod != NULL);

module_state *state = PyModule_GetState(mod);

assert(state != NULL);

return state;

}

static int

traverse_module_state(module_state *state, visitproc visit, void *arg)

{

/* heap types */

Py_VISIT(state->ChannelIDType);

/* exceptions */

Py_VISIT(state->ChannelError);

Py_VISIT(state->ChannelNotFoundError);

Py_VISIT(state->ChannelClosedError);

Py_VISIT(state->ChannelEmptyError);

Py_VISIT(state->ChannelNotEmptyError);

return 0;

}

static int

clear_module_state(module_state *state)

{

/* heap types */

if (state->ChannelIDType != NULL) {

(void)_PyCrossInterpreterData_UnregisterClass(state->ChannelIDType);

}

Py_CLEAR(state->ChannelIDType);

/* exceptions */

Py_CLEAR(state->ChannelError);

Py_CLEAR(state->ChannelNotFoundError);

Py_CLEAR(state->ChannelClosedError);

Py_CLEAR(state->ChannelEmptyError);

Py_CLEAR(state->ChannelNotEmptyError);

return 0;

}

#define CHANNEL_SEND 1

#define CHANNEL_BOTH 0

#define CHANNEL_RECV -1

#define ERR_CHANNEL_NOT_FOUND -2

#define ERR_CHANNEL_CLOSED -3

#define ERR_CHANNEL_INTERP_CLOSED -4

#define ERR_CHANNEL_EMPTY -5

#define ERR_CHANNEL_NOT_EMPTY -6

#define ERR_CHANNEL_MUTEX_INIT -7

#define ERR_CHANNELS_MUTEX_INIT -8

#define ERR_NO_NEXT_CHANNEL_ID -9

static int

exceptions_init(PyObject *mod)

{

module_state *state = get_module_state(mod);

if (state == NULL) {

return -1;

}

#define ADD(NAME, BASE) \

do { \

assert(state->NAME == NULL); \

state->NAME = ADD_NEW_EXCEPTION(mod, NAME, BASE); \

if (state->NAME == NULL) { \

return -1; \

} \

} while (0)

// A channel-related operation failed.

ADD(ChannelError, PyExc_RuntimeError);

// An operation tried to use a channel that doesn't exist.

ADD(ChannelNotFoundError, state->ChannelError);

// An operation tried to use a closed channel.

ADD(ChannelClosedError, state->ChannelError);

// An operation tried to pop from an empty channel.

ADD(ChannelEmptyError, state->ChannelError);

// An operation tried to close a non-empty channel.

ADD(ChannelNotEmptyError, state->ChannelError);

#undef ADD

return 0;

}

static int

handle_channel_error(int err, PyObject *mod, int64_t cid)

{

if (err == 0) {

assert(!PyErr_Occurred());

return 0;

}

assert(err < 0);

module_state *state = get_module_state(mod);

assert(state != NULL);

if (err == ERR_CHANNEL_NOT_FOUND) {

PyErr_Format(state->ChannelNotFoundError,

"channel %" PRId64 " not found", cid);

}

else if (err == ERR_CHANNEL_CLOSED) {

PyErr_Format(state->ChannelClosedError,

"channel %" PRId64 " is closed", cid);

}

else if (err == ERR_CHANNEL_INTERP_CLOSED) {

PyErr_Format(state->ChannelClosedError,

"channel %" PRId64 " is already closed", cid);

}

else if (err == ERR_CHANNEL_EMPTY) {

PyErr_Format(state->ChannelEmptyError,

"channel %" PRId64 " is empty", cid);

}

else if (err == ERR_CHANNEL_NOT_EMPTY) {

PyErr_Format(state->ChannelNotEmptyError,

"channel %" PRId64 " may not be closed "

"if not empty (try force=True)",

cid);

}

else if (err == ERR_CHANNEL_MUTEX_INIT) {

PyErr_SetString(state->ChannelError,

"can't initialize mutex for new channel");

}

else if (err == ERR_CHANNELS_MUTEX_INIT) {

PyErr_SetString(state->ChannelError,

"can't initialize mutex for channel management");

}

else if (err == ERR_NO_NEXT_CHANNEL_ID) {

PyErr_SetString(state->ChannelError,

"failed to get a channel ID");

}

else {

assert(PyErr_Occurred());

}

return 1;

}

struct _channelitem;

typedef struct _channelitem {

_PyCrossInterpreterData *data;

struct _channelitem *next;

} _channelitem;

static _channelitem *

_channelitem_new(void)

{

_channelitem *item = GLOBAL_MALLOC(_channelitem);

if (item == NULL) {

PyErr_NoMemory();

return NULL;

}

item->data = NULL;

item->next = NULL;

return item;

}

static void

_channelitem_clear(_channelitem *item)

{

if (item->data != NULL) {

(void)_release_xid_data(item->data, 1);

// It was allocated in _channel_send().

GLOBAL_FREE(item->data);

item->data = NULL;

}

item->next = NULL;

}

static void

_channelitem_free(_channelitem *item)

{

_channelitem_clear(item);

GLOBAL_FREE(item);

}

static void

_channelitem_free_all(_channelitem *item)

{

while (item != NULL) {

_channelitem *last = item;

item = item->next;

_channelitem_free(last);

}

}

static _PyCrossInterpreterData *

_channelitem_popped(_channelitem *item)

{

_PyCrossInterpreterData *data = item->data;

item->data = NULL;

_channelitem_free(item);

return data;

}

typedef struct _channelqueue {

} _channelqueue;

static _channelqueue *

_channelqueue_new(void)

{

_channelqueue *queue = GLOBAL_MALLOC(_channelqueue);

if (queue == NULL) {

PyErr_NoMemory();

return NULL;

}

queue->count = 0;

queue->first = NULL;

queue->last = NULL;

return queue;

}

static void

_channelqueue_clear(_channelqueue *queue)

{

_channelitem_free_all(queue->first);

queue->count = 0;

queue->first = NULL;

queue->last = NULL;

}

static void

_channelqueue_free(_channelqueue *queue)

{

_channelqueue_clear(queue);

GLOBAL_FREE(queue);

}

static int

_channelqueue_put(_channelqueue *queue, _PyCrossInterpreterData *data)

{

_channelitem *item = _channelitem_new();

if (item == NULL) {

return -1;

}

item->data = data;

queue->count += 1;

if (queue->first == NULL) {

queue->first = item;

}

else {

queue->last->next = item;

}

queue->last = item;

return 0;

}

static _PyCrossInterpreterData *

_channelqueue_get(_channelqueue *queue)

{

_channelitem *item = queue->first;

if (item == NULL) {

return NULL;

}

queue->first = item->next;

if (queue->last == item) {

queue->last = NULL;

}

queue->count -= 1;

return _channelitem_popped(item);

}

static void

_channelqueue_drop_interpreter(_channelqueue *queue, int64_t interp)

{

_channelitem *prev = NULL;

_channelitem *next = queue->first;

while (next != NULL) {

_channelitem *item = next;

next = item->next;

if (item->data->interp == interp) {

if (prev == NULL) {

queue->first = item->next;

}

else {

prev->next = item->next;

}

_channelitem_free(item);

queue->count -= 1;

}

else {

prev = item;

}

}

}

struct _channelend;

typedef struct _channelend {

} _channelend;

static _channelend *

_channelend_new(int64_t interp)

{

_channelend *end = GLOBAL_MALLOC(_channelend);

if (end == NULL) {

PyErr_NoMemory();

return NULL;

}

end->next = NULL;

end->interp = interp;

end->open = 1;

return end;

}

static void

_channelend_free(_channelend *end)

{

GLOBAL_FREE(end);

}

static void

_channelend_free_all(_channelend *end)

{

while (end != NULL) {

_channelend *last = end;

end = end->next;

_channelend_free(last);

}

}

static _channelend *

_channelend_find(_channelend *first, int64_t interp, _channelend **pprev)

{

_channelend *prev = NULL;

_channelend *end = first;

while (end != NULL) {

if (end->interp == interp) {

break;

}

prev = end;

end = end->next;

}

if (pprev != NULL) {

*pprev = prev;

}

return end;

}

typedef struct _channelassociations {

} _channelends;

static _channelends *

_channelends_new(void)

{

_channelends *ends = GLOBAL_MALLOC(_channelends);

if (ends== NULL) {

return NULL;

}

ends->numsendopen = 0;

ends->numrecvopen = 0;

ends->send = NULL;

ends->recv = NULL;

return ends;

}

static void

_channelends_clear(_channelends *ends)

{

_channelend_free_all(ends->send);

ends->send = NULL;

ends->numsendopen = 0;

_channelend_free_all(ends->recv);

ends->recv = NULL;

ends->numrecvopen = 0;

}

static void

_channelends_free(_channelends *ends)

{

_channelends_clear(ends);

GLOBAL_FREE(ends);

}

static _channelend *

_channelends_add(_channelends *ends, _channelend *prev, int64_t interp,

int send)

{

_channelend *end = _channelend_new(interp);

if (end == NULL) {

return NULL;

}

if (prev == NULL) {

if (send) {

ends->send = end;

}

else {

ends->recv = end;

}

}

else {

prev->next = end;

}

if (send) {

ends->numsendopen += 1;

}

else {

ends->numrecvopen += 1;

}

return end;

}

static int

_channelends_associate(_channelends *ends, int64_t interp, int send)

{

_channelend *prev;

_channelend *end = _channelend_find(send ? ends->send : ends->recv,

interp, &prev);

if (end != NULL) {

if (!end->open) {

return ERR_CHANNEL_CLOSED;

}

// already associated

return 0;

}

if (_channelends_add(ends, prev, interp, send) == NULL) {

return -1;

}

return 0;

}

static int

_channelends_is_open(_channelends *ends)

{

if (ends->numsendopen != 0 || ends->numrecvopen != 0) {

return 1;

}

if (ends->send == NULL && ends->recv == NULL) {

return 1;

}

return 0;

}

static void

_channelends_close_end(_channelends *ends, _channelend *end, int send)

{

end->open = 0;

if (send) {

ends->numsendopen -= 1;

}

else {

ends->numrecvopen -= 1;

}

}

static int

_channelends_close_interpreter(_channelends *ends, int64_t interp, int which)

{

_channelend *prev;

_channelend *end;

if (which >= 0) { // send/both

end = _channelend_find(ends->send, interp, &prev);

if (end == NULL) {

// never associated so add it

end = _channelends_add(ends, prev, interp, 1);

if (end == NULL) {

return -1;

}

}

_channelends_close_end(ends, end, 1);

}

if (which <= 0) { // recv/both

end = _channelend_find(ends->recv, interp, &prev);

if (end == NULL) {

// never associated so add it

end = _channelends_add(ends, prev, interp, 0);

if (end == NULL) {

return -1;

}

}

_channelends_close_end(ends, end, 0);

}

return 0;

}

static void

_channelends_drop_interpreter(_channelends *ends, int64_t interp)

{

_channelend *end;

end = _channelend_find(ends->send, interp, NULL);

if (end != NULL) {

_channelends_close_end(ends, end, 1);

}

end = _channelend_find(ends->recv, interp, NULL);

if (end != NULL) {

_channelends_close_end(ends, end, 0);

}

}

static void

_channelends_close_all(_channelends *ends, int which, int force)

{

// XXX Handle the ends.

// XXX Handle force is True.

// Ensure all the "send"-associated interpreters are closed.

_channelend *end;

for (end = ends->send; end != NULL; end = end->next) {

_channelends_close_end(ends, end, 1);

}

// Ensure all the "recv"-associated interpreters are closed.

for (end = ends->recv; end != NULL; end = end->next) {

_channelends_close_end(ends, end, 0);

}

}

struct _channel;

struct _channel_closing;

static void _channel_clear_closing(struct _channel *);

static void _channel_finish_closing(struct _channel *);

typedef struct _channel {

} _PyChannelState;

static _PyChannelState *

_channel_new(PyThread_type_lock mutex)

{

_PyChannelState *chan = GLOBAL_MALLOC(_PyChannelState);

if (chan == NULL) {

return NULL;

}

chan->mutex = mutex;

chan->queue = _channelqueue_new();

if (chan->queue == NULL) {

GLOBAL_FREE(chan);

return NULL;

}

chan->ends = _channelends_new();

if (chan->ends == NULL) {

_channelqueue_free(chan->queue);

GLOBAL_FREE(chan);

return NULL;

}

chan->open = 1;

chan->closing = NULL;

return chan;

}

static void

_channel_free(_PyChannelState *chan)

{

_channel_clear_closing(chan);

PyThread_acquire_lock(chan->mutex, WAIT_LOCK);

_channelqueue_free(chan->queue);

_channelends_free(chan->ends);

PyThread_release_lock(chan->mutex);

PyThread_free_lock(chan->mutex);

GLOBAL_FREE(chan);

}

static int

_channel_add(_PyChannelState *chan, int64_t interp,

_PyCrossInterpreterData *data)

{

int res = -1;

PyThread_acquire_lock(chan->mutex, WAIT_LOCK);

if (!chan->open) {

res = ERR_CHANNEL_CLOSED;

goto done;

}

if (_channelends_associate(chan->ends, interp, 1) != 0) {

res = ERR_CHANNEL_INTERP_CLOSED;

goto done;

}

if (_channelqueue_put(chan->queue, data) != 0) {

goto done;

}

res = 0;

done:

PyThread_release_lock(chan->mutex);

return res;

}

static int

_channel_next(_PyChannelState *chan, int64_t interp,

_PyCrossInterpreterData **res)

{

int err = 0;

PyThread_acquire_lock(chan->mutex, WAIT_LOCK);

if (!chan->open) {

err = ERR_CHANNEL_CLOSED;

goto done;

}

if (_channelends_associate(chan->ends, interp, 0) != 0) {

err = ERR_CHANNEL_INTERP_CLOSED;

goto done;

}

_PyCrossInterpreterData *data = _channelqueue_get(chan->queue);

if (data == NULL && !PyErr_Occurred() && chan->closing != NULL) {

chan->open = 0;

}

*res = data;

done:

PyThread_release_lock(chan->mutex);

if (chan->queue->count == 0) {

_channel_finish_closing(chan);

}

return err;

}

static int

_channel_close_interpreter(_PyChannelState *chan, int64_t interp, int end)

{

PyThread_acquire_lock(chan->mutex, WAIT_LOCK);

int res = -1;

if (!chan->open) {

res = ERR_CHANNEL_CLOSED;

goto done;

}

if (_channelends_close_interpreter(chan->ends, interp, end) != 0) {

goto done;

}

chan->open = _channelends_is_open(chan->ends);

res = 0;

done:

PyThread_release_lock(chan->mutex);

return res;

}

static void

_channel_drop_interpreter(_PyChannelState *chan, int64_t interp)

{

PyThread_acquire_lock(chan->mutex, WAIT_LOCK);

_channelqueue_drop_interpreter(chan->queue, interp);

_channelends_drop_interpreter(chan->ends, interp);

chan->open = _channelends_is_open(chan->ends);

PyThread_release_lock(chan->mutex);

}

static int

_channel_close_all(_PyChannelState *chan, int end, int force)

{

int res = -1;

PyThread_acquire_lock(chan->mutex, WAIT_LOCK);

if (!chan->open) {

res = ERR_CHANNEL_CLOSED;

goto done;

}

if (!force && chan->queue->count > 0) {

res = ERR_CHANNEL_NOT_EMPTY;

goto done;

}

chan->open = 0;

// We *could* also just leave these in place, since we've marked

// the channel as closed already.

_channelends_close_all(chan->ends, end, force);

res = 0;

done:

PyThread_release_lock(chan->mutex);

return res;

}

struct _channelref;

typedef struct _channelref {

} _channelref;

static _channelref *

_channelref_new(int64_t id, _PyChannelState *chan)

{

_channelref *ref = GLOBAL_MALLOC(_channelref);

if (ref == NULL) {

return NULL;

}

ref->id = id;

ref->chan = chan;

ref->next = NULL;

ref->objcount = 0;

return ref;

}

static void

_channelref_free(_channelref *ref)

{

if (ref->chan != NULL) {

_channel_clear_closing(ref->chan);

}

//_channelref_clear(ref);

GLOBAL_FREE(ref);

}

static _channelref *

_channelref_find(_channelref *first, int64_t id, _channelref **pprev)

{

_channelref *prev = NULL;

_channelref *ref = first;

while (ref != NULL) {

if (ref->id == id) {

break;

}

prev = ref;

ref = ref->next;

}

if (pprev != NULL) {

*pprev = prev;

}

return ref;

}

typedef struct _channels {

} _channels;

static void

_channels_init(_channels *channels, PyThread_type_lock mutex)

{

channels->mutex = mutex;

channels->head = NULL;

channels->numopen = 0;

channels->next_id = 0;

}

static void

_channels_fini(_channels *channels)

{

assert(channels->numopen == 0);

assert(channels->head == NULL);

if (channels->mutex != NULL) {

PyThread_free_lock(channels->mutex);

channels->mutex = NULL;

}

}

static int64_t

_channels_next_id(_channels *channels) // needs lock

{

int64_t id = channels->next_id;

if (id < 0) {

/* overflow */