limingjie/OmniSocketGo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: 对接Python，暴露接口

Browse Source
This commit is contained in:
Mock
2026-03-30 22:48:36 +08:00
parent 24467c04c0
commit d678bfc326
22 changed files with 1311 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
python/omnisocket/__init__.py Normal file
View File

@@ -0,0 +1,44 @@
try:
from ._omnisocket import (
MSG_TYPE_BINARY,
MSG_TYPE_ERROR,
MSG_TYPE_FILE,
MSG_TYPE_REGISTER,
MSG_TYPE_TEXT,
Session,
)
except ImportError as exc:
raise ImportError(
"omnisocket extension is not built; run `make python-ext` on a Linux host first"
) from exc
CONTROL_DEFAULTS = {
"nodelay": 1,
"interval_ms": 5,
"resend": 2,
"nc": 1,
"sndwnd": 32,
"rcvwnd": 32,
"mtu": 1400,
}
VIDEO_DEFAULTS = {
"nodelay": 1,
"interval_ms": 10,
"resend": 2,
"nc": 1,
"sndwnd": 256,
"rcvwnd": 256,
"mtu": 1400,
}
__all__ = [
"CONTROL_DEFAULTS",
"VIDEO_DEFAULTS",
"MSG_TYPE_BINARY",
"MSG_TYPE_ERROR",
"MSG_TYPE_FILE",
"MSG_TYPE_REGISTER",
"MSG_TYPE_TEXT",
"Session",
]

336
python/omnisocket/_omnisocket.c Normal file
View File

@@ -0,0 +1,336 @@
#define PY_SSIZE_T_CLEAN
#include <Python.h>
#include "omnisocket_client.h"
typedef struct PyOmniSession {
PyObject_HEAD
omnisocket_session_t session;
} PyOmniSession;
static const char *PyOmniSession_recv_doc =
"recv(timeout_ms=-1) -> (from_peer, msg_type, payload) | None";
static const char *PyOmniSession_recv_into_doc =
"recv_into(buffer, timeout_ms=-1) -> dict | None\n"
"\n"
"The writable buffer must be large enough for the full message body.\n"
"If it is too small, BufferError reports the required size but the\n"
"current frame has already been consumed and is lost.";
static PyObject *PyOmniSession_new(PyTypeObject *type, PyObject *args, PyObject *kwargs) {
PyOmniSession *self;
(void) args;
(void) kwargs;
self = (PyOmniSession *) type->tp_alloc(type, 0);
if (self == NULL) {
return NULL;
}
if (omnisocket_session_init(&self->session) != 0) {
type->tp_free((PyObject *) self);
return PyErr_SetFromErrno(PyExc_OSError);
}
return (PyObject *) self;
}
static void PyOmniSession_dealloc(PyOmniSession *self) {
omnisocket_session_destroy(&self->session);
Py_TYPE(self)->tp_free((PyObject *) self);
}
static PyObject *PyOmniSession_connect(PyOmniSession *self, PyObject *args, PyObject *kwargs) {
const char *server_addr;
const char *peer_id;
const char *relay_via = "";
const char *bind_ip = "";
const char *bind_device = "";
int nodelay = KCP_DEFAULT_NODELAY;
int interval_ms = KCP_DEFAULT_INTERVAL_MS;
int resend = KCP_DEFAULT_RESEND;
int nc = KCP_DEFAULT_NC;
int sndwnd = KCP_DEFAULT_SND_WND;
int rcvwnd = KCP_DEFAULT_RCV_WND;
int mtu = KCP_DEFAULT_MTU;
int stats_interval_ms = KCP_DEFAULT_STATS_INTERVAL_MS;
kcp_conn_options_t options;
int rc;
static char *kwlist[] = {
"server_addr",
"peer_id",
"relay_via",
"bind_ip",
"bind_device",
"nodelay",
"interval_ms",
"resend",
"nc",
"sndwnd",
"rcvwnd",
"mtu",
"stats_interval_ms",
NULL
};
if (!PyArg_ParseTupleAndKeywords(
args,
kwargs,
"ss|sssiiiiiiii",
kwlist,
&server_addr,
&peer_id,
&relay_via,
&bind_ip,
&bind_device,
&nodelay,
&interval_ms,
&resend,
&nc,
&sndwnd,
&rcvwnd,
&mtu,
&stats_interval_ms)) {
return NULL;
}
kcp_conn_options_init(&options);
options.nodelay = nodelay;
options.interval_ms = interval_ms;
options.resend = resend;
options.nc = nc;
options.sndwnd = sndwnd;
options.rcvwnd = rcvwnd;
options.mtu = mtu;
Py_BEGIN_ALLOW_THREADS
rc = omnisocket_session_connect(
&self->session,
server_addr,
relay_via,
peer_id,
bind_ip,
bind_device,
&options,
stats_interval_ms
);
Py_END_ALLOW_THREADS
if (rc != 0) {
return PyErr_SetFromErrno(PyExc_OSError);
}
Py_RETURN_NONE;
}
static PyObject *PyOmniSession_close(PyOmniSession *self, PyObject *Py_UNUSED(ignored)) {
int rc;
Py_BEGIN_ALLOW_THREADS
rc = omnisocket_session_close(&self->session);
Py_END_ALLOW_THREADS
if (rc != 0) {
return PyErr_SetFromErrno(PyExc_OSError);
}
Py_RETURN_NONE;
}
static PyObject *PyOmniSession_send(PyOmniSession *self, PyObject *args, PyObject *kwargs) {
const char *to;
Py_buffer payload;
int rc;
static char *kwlist[] = {"to", "data", NULL};
memset(&payload, 0, sizeof(payload));
if (!PyArg_ParseTupleAndKeywords(args, kwargs, "sy*", kwlist, &to, &payload)) {
return NULL;
}
Py_BEGIN_ALLOW_THREADS
rc = omnisocket_session_send(&self->session, to, payload.buf, (size_t) payload.len);
Py_END_ALLOW_THREADS
PyBuffer_Release(&payload);
if (rc != 0) {
return PyErr_SetFromErrno(PyExc_OSError);
}
Py_RETURN_NONE;
}
static PyObject *PyOmniSession_recv(PyOmniSession *self, PyObject *args, PyObject *kwargs) {
int timeout_ms = -1;
int rc;
message_t msg;
PyObject *body = NULL;
PyObject *result = NULL;
static char *kwlist[] = {"timeout_ms", NULL};
if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|i", kwlist, &timeout_ms)) {
return NULL;
}
protocol_message_init(&msg);
Py_BEGIN_ALLOW_THREADS
rc = omnisocket_session_recv(&self->session, &msg, timeout_ms);
Py_END_ALLOW_THREADS
if (rc == 1) {
protocol_message_clear(&msg);
Py_RETURN_NONE;
}
if (rc != 0) {
protocol_message_clear(&msg);
return PyErr_SetFromErrno(PyExc_OSError);
}
body = PyBytes_FromStringAndSize((const char *) msg.body, (Py_ssize_t) msg.body_len);
if (body == NULL) {
protocol_message_clear(&msg);
return NULL;
}
result = Py_BuildValue("(siO)", msg.from, (int) msg.type, body);
Py_DECREF(body);
protocol_message_clear(&msg);
return result;
}
static PyObject *PyOmniSession_recv_into(PyOmniSession *self, PyObject *args, PyObject *kwargs) {
PyObject *buffer_obj;
Py_buffer view;
int timeout_ms = -1;
int rc;
kcp_client_recv_meta_t meta;
PyObject *result = NULL;
static char *kwlist[] = {"buffer", "timeout_ms", NULL};
memset(&view, 0, sizeof(view));
memset(&meta, 0, sizeof(meta));
if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|i", kwlist, &buffer_obj, &timeout_ms)) {
return NULL;
}
if (PyObject_GetBuffer(buffer_obj, &view, PyBUF_WRITABLE) != 0) {
return NULL;
}
Py_BEGIN_ALLOW_THREADS
rc = omnisocket_session_recv_into(&self->session, view.buf, (size_t) view.len, &meta, timeout_ms);
Py_END_ALLOW_THREADS
PyBuffer_Release(&view);
if (rc == 1) {
Py_RETURN_NONE;
}
if (rc == 2) {
PyErr_Format(
PyExc_BufferError,
"buffer too small: need %zu bytes; current frame was already consumed and dropped",
meta.body_len
);
return NULL;
}
if (rc != 0) {
return PyErr_SetFromErrno(PyExc_OSError);
}
result = Py_BuildValue(
"{s:s,s:s,s:s,s:i,s:K,s:K}",
"from",
meta.from,
"to",
meta.to,
"file_name",
meta.file_name,
"msg_type",
(int) meta.type,
"message_id",
(unsigned long long) meta.id,
"body_len",
(unsigned long long) meta.body_len
);
return result;
}
static PyObject *PyOmniSession_stats(PyOmniSession *self, PyObject *Py_UNUSED(ignored)) {
omnisocket_session_stats_t stats;
memset(&stats, 0, sizeof(stats));
omnisocket_session_stats_snapshot(&self->session, &stats);
return Py_BuildValue(
"{s:K,s:K,s:K,s:K,s:K,s:K,s:K,s:i}",
"send_calls",
(unsigned long long) stats.send_calls,
"send_bytes",
(unsigned long long) stats.send_bytes,
"send_errors",
(unsigned long long) stats.send_errors,
"recv_calls",
(unsigned long long) stats.recv_calls,
"recv_bytes",
(unsigned long long) stats.recv_bytes,
"recv_timeouts",
(unsigned long long) stats.recv_timeouts,
"recv_errors",
(unsigned long long) stats.recv_errors,
"connected",
stats.connected
);
}
static PyMethodDef PyOmniSession_methods[] = {
{"connect", (PyCFunction) PyOmniSession_connect, METH_VARARGS | METH_KEYWORDS, NULL},
{"close", (PyCFunction) PyOmniSession_close, METH_NOARGS, NULL},
{"send", (PyCFunction) PyOmniSession_send, METH_VARARGS | METH_KEYWORDS, NULL},
{"recv", (PyCFunction) PyOmniSession_recv, METH_VARARGS | METH_KEYWORDS, PyOmniSession_recv_doc},
{"recv_into", (PyCFunction) PyOmniSession_recv_into, METH_VARARGS | METH_KEYWORDS, PyOmniSession_recv_into_doc},
{"stats", (PyCFunction) PyOmniSession_stats, METH_NOARGS, NULL},
{NULL, NULL, 0, NULL}
};
static PyTypeObject PyOmniSessionType = {
PyVarObject_HEAD_INIT(NULL, 0)
};
static PyModuleDef omnisocket_module = {
PyModuleDef_HEAD_INIT,
.m_name = "_omnisocket",
.m_size = -1,
};
PyMODINIT_FUNC PyInit__omnisocket(void) {
PyObject *module;
PyOmniSessionType.tp_name = "omnisocket.Session";
PyOmniSessionType.tp_basicsize = sizeof(PyOmniSession);
PyOmniSessionType.tp_flags = Py_TPFLAGS_DEFAULT;
PyOmniSessionType.tp_new = PyOmniSession_new;
PyOmniSessionType.tp_dealloc = (destructor) PyOmniSession_dealloc;
PyOmniSessionType.tp_methods = PyOmniSession_methods;
if (PyType_Ready(&PyOmniSessionType) < 0) {
return NULL;
}
module = PyModule_Create(&omnisocket_module);
if (module == NULL) {
return NULL;
}
Py_INCREF(&PyOmniSessionType);
if (PyModule_AddObject(module, "Session", (PyObject *) &PyOmniSessionType) != 0) {
Py_DECREF(&PyOmniSessionType);
Py_DECREF(module);
return NULL;
}
if (PyModule_AddIntConstant(module, "MSG_TYPE_TEXT", MSG_TYPE_TEXT) != 0 ||
PyModule_AddIntConstant(module, "MSG_TYPE_FILE", MSG_TYPE_FILE) != 0 ||
PyModule_AddIntConstant(module, "MSG_TYPE_REGISTER", MSG_TYPE_REGISTER) != 0 ||
PyModule_AddIntConstant(module, "MSG_TYPE_ERROR", MSG_TYPE_ERROR) != 0 ||
PyModule_AddIntConstant(module, "MSG_TYPE_BINARY", MSG_TYPE_BINARY) != 0) {
Py_DECREF(module);
return NULL;
}
return module;
}

248
python/omnisocket/omnisocket_client.c Normal file
View File

@@ -0,0 +1,248 @@
#include "omnisocket_client.h"
int omnisocket_session_init(omnisocket_session_t *session) {
int rc;
if (session == NULL) {
errno = EINVAL;
return -1;
}
memset(session, 0, sizeof(*session));
rc = pthread_mutex_init(&session->mutex, NULL);
if (rc != 0) {
errno = rc;
return -1;
}
rc = pthread_cond_init(&session->idle_cond, NULL);
if (rc != 0) {
pthread_mutex_destroy(&session->mutex);
errno = rc;
return -1;
}
return 0;
}
void omnisocket_session_destroy(omnisocket_session_t *session) {
if (session == NULL) {
return;
}
(void) omnisocket_session_close(session);
pthread_cond_destroy(&session->idle_cond);
pthread_mutex_destroy(&session->mutex);
}
static int omnisocket_session_begin_client_op(omnisocket_session_t *session, kcp_client_t **out_client) {
if (session == NULL || out_client == NULL) {
errno = EINVAL;
return -1;
}
pthread_mutex_lock(&session->mutex);
if (session->closing) {
pthread_mutex_unlock(&session->mutex);
errno = ECANCELED;
return -1;
}
if (session->client == NULL) {
pthread_mutex_unlock(&session->mutex);
errno = ENOTCONN;
return -1;
}
*out_client = session->client;
session->active_ops += 1;
pthread_mutex_unlock(&session->mutex);
return 0;
}
int omnisocket_session_connect(
omnisocket_session_t *session,
const char *server_addr,
const char *relay_via,
const char *peer_id,
const char *bind_ip,
const char *bind_device,
const kcp_conn_options_t *options,
int stats_interval_ms
) {
kcp_client_t *client;
if (session == NULL || server_addr == NULL || peer_id == NULL) {
errno = EINVAL;
return -1;
}
pthread_mutex_lock(&session->mutex);
while (session->closing) {
pthread_cond_wait(&session->idle_cond, &session->mutex);
}
if (session->client != NULL) {
pthread_mutex_unlock(&session->mutex);
errno = EISCONN;
return -1;
}
client = kcp_client_dial_with_options(
server_addr,
relay_via,
peer_id,
bind_ip,
bind_device,
options,
NULL,
NULL,
NULL,
stats_interval_ms
);
if (client == NULL) {
pthread_mutex_unlock(&session->mutex);
return -1;
}
session->client = client;
session->stats.connected = 1;
pthread_mutex_unlock(&session->mutex);
return 0;
}
int omnisocket_session_close(omnisocket_session_t *session) {
kcp_client_t *client;
if (session == NULL) {
errno = EINVAL;
return -1;
}
pthread_mutex_lock(&session->mutex);
while (session->closing) {
pthread_cond_wait(&session->idle_cond, &session->mutex);
}
client = session->client;
if (client != NULL) {
session->closing = 1;
session->client = NULL;
}
session->stats.connected = 0;
pthread_mutex_unlock(&session->mutex);
if (client != NULL) {
kcp_client_close(client);
pthread_mutex_lock(&session->mutex);
while (session->active_ops > 0) {
pthread_cond_wait(&session->idle_cond, &session->mutex);
}
pthread_mutex_unlock(&session->mutex);
kcp_client_free(client);
pthread_mutex_lock(&session->mutex);
session->closing = 0;
pthread_cond_broadcast(&session->idle_cond);
pthread_mutex_unlock(&session->mutex);
}
return 0;
}
int omnisocket_session_send(omnisocket_session_t *session, const char *to, const void *data, size_t data_len) {
kcp_client_t *client;
int rc;
if (session == NULL || to == NULL || (data == NULL && data_len > 0)) {
errno = EINVAL;
return -1;
}
if (omnisocket_session_begin_client_op(session, &client) != 0) {
return -1;
}
rc = kcp_client_send_binary(client, to, data, data_len);
pthread_mutex_lock(&session->mutex);
if (rc == 0) {
session->stats.send_calls += 1;
session->stats.send_bytes += (uint64_t) data_len;
} else {
session->stats.send_errors += 1;
}
if (session->active_ops > 0) {
session->active_ops -= 1;
}
if (session->closing && session->active_ops == 0) {
pthread_cond_broadcast(&session->idle_cond);
}
pthread_mutex_unlock(&session->mutex);
return rc;
}
int omnisocket_session_recv(omnisocket_session_t *session, message_t *out_msg, int timeout_ms) {
kcp_client_t *client;
int rc;
if (session == NULL || out_msg == NULL) {
errno = EINVAL;
return -1;
}
if (omnisocket_session_begin_client_op(session, &client) != 0) {
return -1;
}
rc = kcp_client_receive_timed(client, out_msg, timeout_ms);
pthread_mutex_lock(&session->mutex);
if (rc == 0) {
session->stats.recv_calls += 1;
session->stats.recv_bytes += (uint64_t) out_msg->body_len;
} else if (rc == 1) {
session->stats.recv_timeouts += 1;
} else {
session->stats.recv_errors += 1;
}
if (session->active_ops > 0) {
session->active_ops -= 1;
}
if (session->closing && session->active_ops == 0) {
pthread_cond_broadcast(&session->idle_cond);
}
pthread_mutex_unlock(&session->mutex);
return rc;
}
int omnisocket_session_recv_into(
omnisocket_session_t *session,
void *buffer,
size_t buffer_len,
kcp_client_recv_meta_t *out_meta,
int timeout_ms
) {
kcp_client_t *client;
int rc;
if (session == NULL || out_meta == NULL || (buffer == NULL && buffer_len > 0)) {
errno = EINVAL;
return -1;
}
if (omnisocket_session_begin_client_op(session, &client) != 0) {
return -1;
}
rc = kcp_client_receive_binary_into(client, buffer, buffer_len, out_meta, timeout_ms);
pthread_mutex_lock(&session->mutex);
if (rc == 0) {
session->stats.recv_calls += 1;
session->stats.recv_bytes += (uint64_t) out_meta->body_len;
} else if (rc == 1) {
session->stats.recv_timeouts += 1;
} else {
session->stats.recv_errors += 1;
}
if (session->active_ops > 0) {
session->active_ops -= 1;
}
if (session->closing && session->active_ops == 0) {
pthread_cond_broadcast(&session->idle_cond);
}
pthread_mutex_unlock(&session->mutex);
return rc;
}
void omnisocket_session_stats_snapshot(omnisocket_session_t *session, omnisocket_session_stats_t *out_stats) {
if (session == NULL || out_stats == NULL) {
return;
}
pthread_mutex_lock(&session->mutex);
*out_stats = session->stats;
pthread_mutex_unlock(&session->mutex);
}

51
python/omnisocket/omnisocket_client.h Normal file
View File

@@ -0,0 +1,51 @@
#ifndef OMNISOCKET_PY_CLIENT_H
#define OMNISOCKET_PY_CLIENT_H
#include "peer_kcp_client.h"
typedef struct omnisocket_session_stats {
uint64_t send_calls;
uint64_t send_bytes;
uint64_t send_errors;
uint64_t recv_calls;
uint64_t recv_bytes;
uint64_t recv_timeouts;
uint64_t recv_errors;
int connected;
} omnisocket_session_stats_t;
typedef struct omnisocket_session {
pthread_mutex_t mutex;
pthread_cond_t idle_cond;
kcp_client_t *client;
size_t active_ops;
int closing;
omnisocket_session_stats_t stats;
} omnisocket_session_t;
int omnisocket_session_init(omnisocket_session_t *session);
void omnisocket_session_destroy(omnisocket_session_t *session);
int omnisocket_session_connect(
omnisocket_session_t *session,
const char *server_addr,
const char *relay_via,
const char *peer_id,
const char *bind_ip,
const char *bind_device,
const kcp_conn_options_t *options,
int stats_interval_ms
);
int omnisocket_session_close(omnisocket_session_t *session);
int omnisocket_session_send(omnisocket_session_t *session, const char *to, const void *data, size_t data_len);
int omnisocket_session_recv(omnisocket_session_t *session, message_t *out_msg, int timeout_ms);
int omnisocket_session_recv_into(
omnisocket_session_t *session,
void *buffer,
size_t buffer_len,
kcp_client_recv_meta_t *out_meta,
int timeout_ms
);
void omnisocket_session_stats_snapshot(omnisocket_session_t *session, omnisocket_session_stats_t *out_stats);
#endif

58
python/setup.py Normal file
View File

@@ -0,0 +1,58 @@
from pathlib import Path
import sys
from setuptools import Extension, setup
ROOT = Path(__file__).resolve().parent.parent
PY_ROOT = Path(__file__).resolve().parent
if sys.platform != "linux":
raise RuntimeError("omnisocket Python extension can only be built on Linux")
COMMON_SOURCES = [
ROOT / "src" / "omni_common.c",
ROOT / "src" / "protocol.c",
ROOT / "src" / "latencylog.c",
ROOT / "src" / "tx_timestamp_debug.c",
ROOT / "src" / "kcp_packet_debug.c",
ROOT / "src" / "kcp_session_stats.c",
ROOT / "src" / "linux_timestamping.c",
ROOT / "src" / "interactive.c",
ROOT / "src" / "transport_udp.c",
ROOT / "src" / "transport_kcp.c",
ROOT / "src" / "server_udp_relay.c",
ROOT / "src" / "server_udp_hub.c",
ROOT / "src" / "server_kcp_hub.c",
ROOT / "src" / "peer_udp_client.c",
ROOT / "src" / "peer_kcp_client.c",
ROOT / "third_party" / "cjson" / "cJSON.c",
ROOT / "third_party" / "kcp" / "ikcp.c",
]
setup(
name="omnisocket",
version="0.1.0",
packages=["omnisocket"],
ext_modules=[
Extension(
"omnisocket._omnisocket",
sources=[
str(PY_ROOT / "omnisocket" / "_omnisocket.c"),
str(PY_ROOT / "omnisocket" / "omnisocket_client.c"),
*[str(path) for path in COMMON_SOURCES],
],
include_dirs=[
str(ROOT / "include"),
str(ROOT / "third_party" / "cjson"),
str(ROOT / "third_party" / "kcp"),
str(PY_ROOT / "omnisocket"),
],
define_macros=[("_GNU_SOURCE", None)],
extra_compile_args=["-std=c11", "-O2", "-pthread"],
extra_link_args=["-pthread"],
)
],
)