int flock_open(lua_State *L) {
const char *file_name = lua_tostring(L, 1);
int no_block = luaT_optboolean(L, 2, 0);
int flags = O_CLOEXEC | O_RDWR;
if (!no_block) {
flags |= O_CREAT;
}
int fd = open(file_name, flags, S_IRUSR | S_IWUSR);
if (fd < 0) {
if (errno == ENOENT || errno == EACCES) return 0;
return LUA_HANDLE_ERROR(L, errno);
}
flags = LOCK_EX;
if (no_block) {
flags |= LOCK_NB;
}
int ret = flock(fd, flags);
if (ret < 0) {
close(fd);
if ((flags & LOCK_NB) && (errno == EWOULDBLOCK)) return 0;
return LUA_HANDLE_ERROR(L, errno);
}
int *handle = lua_newuserdata(L, sizeof(int));
*handle = fd;
luaL_getmetatable(L, "ipc.flock");
lua_setmetatable(L, -2);
return 1;
}
int map_open(lua_State *L) {
uint32_t num_threads = lua_tonumber(L, 1);
if (lua_type(L, 2) != LUA_TFUNCTION) return LUA_HANDLE_ERROR_STR(L, "map arg #2 expected a function");
map_thread_t *threads = (map_thread_t *)calloc(num_threads, sizeof(map_thread_t));
int k = lua_gettop(L);
for (uint32_t i = 0; i < num_threads; i++) {
threads[i].rb = ringbuffer_create(MAX_ARG_SIZE);
for (int j = 2; j <= k; j++) {
int ret = rb_save_with_growth(L, j, threads[i].rb);
if (ret) return LUA_HANDLE_ERROR(L, ret);
}
lua_pushinteger(L, i + 1);
int ret = rb_save_with_growth(L, k + 1, threads[i].rb);
if (ret) return LUA_HANDLE_ERROR(L, ret);
lua_pop(L, 1);
ret = pthread_create(&threads[i].thread, NULL, thread_func, &threads[i]);
if (ret) return LUA_HANDLE_ERROR(L, ret);
}
map_t *map = (map_t *)lua_newuserdata(L, sizeof(map_t));
map->num_threads = num_threads;
map->threads = threads;
luaL_getmetatable(L, "ipc.map");
lua_setmetatable(L, -2);
return 1;
}
int cliser_client(lua_State *L) {
#ifdef USE_CUDA
// sometimes cutorch is loaded late, this is a good spot to try and register...
Lcliser_CudaInit(L);
#endif
const char *host;
const char *port;
if (lua_type(L, 1) == LUA_TSTRING) {
host = lua_tostring(L, 1);
port = lua_tostring(L, 2);
} else {
host = DEFAULT_HOST;
port = lua_tostring(L, 1);
}
socklen_t addrlen = sizeof(struct sockaddr);
struct sockaddr addr;
int ret = get_sockaddr(L, host, port, &addr, &addrlen);
if (ret) return ret;
struct timeval tv;
gettimeofday(&tv, NULL);
uint32_t t = tv.tv_sec + DEFAULT_TIMEOUT_SECONDS;
int sock = -1;
while (tv.tv_sec < t) {
ret = socket(PF_INET, SOCK_STREAM, 0);
if (ret <= 0) return LUA_HANDLE_ERROR(L, errno);
sock = ret;
int value = 1;
ret = setsockopt(sock, IPPROTO_TCP, TCP_NODELAY, &value, sizeof(value));
if (ret) break;
ret = connect(sock, &addr, addrlen);
if (!ret) break;
close(sock);
sleep(1);
gettimeofday(&tv, NULL);
}
if (ret) {
close(sock);
return LUA_HANDLE_ERROR(L, errno);
}
addrlen = sizeof(struct sockaddr);
struct sockaddr bind_addr;
ret = getsockname(sock, &bind_addr, &addrlen);
if (ret) {
close(sock);
return LUA_HANDLE_ERROR(L, errno);
}
int use_fastpath = can_use_fastpath(L, sock, ((struct sockaddr_in *)&bind_addr)->sin_addr.s_addr, ((struct sockaddr_in *)&addr)->sin_addr.s_addr);
client_t *client = (client_t *)calloc(1, sizeof(client_t));
client->sock = sock;
client->send_rb = ringbuffer_create(SEND_RECV_SIZE);
client->recv_rb = ringbuffer_create(SEND_RECV_SIZE);
client->ref_count = 1;
client->copy_context.use_fastpath = use_fastpath;
client_t **clientp = (client_t **)lua_newuserdata(L, sizeof(client_t *));
*clientp = client;
luaL_getmetatable(L, "ipc.client");
lua_setmetatable(L, -2);
return 1;
}
static int sock_send_msg(lua_State *L, int index, int sock, ringbuffer_t *rb, copy_context_t *copy_context) {
ringbuffer_push_write_pos(rb);
int ret = rb_save(L, index, rb, 1);
if (ret) return LUA_HANDLE_ERROR(L, ret);
size_t len = ringbuffer_peek(rb);
ringbuffer_pop_write_pos(rb);
if (ret) return ret;
ret = sock_send(sock, &len, sizeof(len), copy_context);
if (ret < 0) return LUA_HANDLE_ERROR(L, errno);
if (ret != sizeof(len)) return LUA_HANDLE_ERROR_STR(L, "failed to send the correct number of bytes");
ret = sock_send(sock, ringbuffer_buf_ptr(rb), len, copy_context);
if (ret < 0) return LUA_HANDLE_ERROR(L, errno);
if ((size_t)ret != len) return LUA_HANDLE_ERROR_STR(L, "failed to send the correct number of bytes");
return 0;
}
int map_join(lua_State *L) {
int rc = 0;
int err_rc = -1;
map_t *map = (map_t *)lua_touserdata(L, 1);
for (uint32_t i = 0; i < map->num_threads; i++) {
if (map->threads[i].rb) {
int ret = pthread_join(map->threads[i].thread, NULL);
if (ret) return LUA_HANDLE_ERROR(L, ret);
if (map->threads[i].ret) {
err_rc = rc;
}
while (ringbuffer_peek(map->threads[i].rb)) {
rb_load(L, map->threads[i].rb);
rc++;
}
ringbuffer_destroy(map->threads[i].rb);
}
}
free(map->threads);
map->threads = NULL;
map->num_threads = 0;
if (err_rc >= 0) {
return LUA_HANDLE_ERROR_STR(L, lua_tostring(L, err_rc - rc));
}
return rc;
}
static int sock_recv_msg_peek(lua_State *L, int sock, ringbuffer_t *rb) {
size_t len;
int ret = recv(sock, &len, sizeof(len), MSG_PEEK | MSG_DONTWAIT);
if (ret < 0) {
if (errno == EAGAIN || errno == EWOULDBLOCK) return 0;
return LUA_HANDLE_ERROR(L, errno);
}
if (ret != sizeof(len)) return 0;
if (len > SEND_RECV_SIZE) return LUA_HANDLE_ERROR_STR(L, "message size is too large");
ret = recv(sock, ringbuffer_buf_ptr(rb), len, MSG_PEEK | MSG_DONTWAIT);
if (ret < 0) {
if (errno == EAGAIN || errno == EWOULDBLOCK) return 0;
return LUA_HANDLE_ERROR(L, errno);
}
if ((size_t)ret != len) return 0;
return ret;
}
static int sock_recv_msg(lua_State *L, int sock, ringbuffer_t *rb, copy_context_t *copy_context) {
size_t len;
int ret = sock_recv(sock, &len, sizeof(len), copy_context);
if (ret < 0) return LUA_HANDLE_ERROR(L, errno);
if (ret != sizeof(len)) return LUA_HANDLE_ERROR_STR(L, "failed to recv the correct number of bytes");
if (len > SEND_RECV_SIZE) return LUA_HANDLE_ERROR_STR(L, "message size is too large");
ret = sock_recv(sock, ringbuffer_buf_ptr(rb), len, copy_context);
if (ret < 0) return LUA_HANDLE_ERROR(L, errno);
if ((size_t)ret != len) return LUA_HANDLE_ERROR_STR(L, "failed to recv the correct number of bytes");
ringbuffer_reset_read_pos(rb);
ringbuffer_push_write_pos(rb);
if (ringbuffer_write(rb, NULL, ret) != (size_t)ret) {
ringbuffer_pop_write_pos(rb);
return LUA_HANDLE_ERROR_STR(L, "failed to write the correct number of bytes into the ringbuffer");
}
ret = rb_load(L, rb);
if (ret < 0) return LUA_HANDLE_ERROR(L, ret);
return ret;
}
int cliser_server(lua_State *L) {
#ifdef USE_CUDA
// sometimes cutorch is loaded late, this is a good spot to try and register...
Lcliser_CudaInit(L);
#endif
const char *host = luaL_optstring(L, 1, DEFAULT_HOST);
int port = luaL_optinteger(L, 2, 0);
char port_str[16];
snprintf(port_str, 16, "%d", port);
struct sockaddr addr;
socklen_t addrlen = sizeof(struct sockaddr);
int ret = get_sockaddr(L, host, port != 0 ? port_str : NULL, &addr, &addrlen);
if (ret) return ret;
ret = socket(PF_INET, SOCK_STREAM, 0);
if (ret <= 0) return LUA_HANDLE_ERROR(L, errno);
int sock = ret;
ret = bind(sock, &addr, addrlen);
if (ret) {
close(sock);
return LUA_HANDLE_ERROR(L, errno);
}
ret = listen(sock, 1024);
if (ret) {
close(sock);
return LUA_HANDLE_ERROR(L, errno);
}
struct sockaddr_in sin;
addrlen = sizeof(struct sockaddr_in);
ret = getsockname(sock, (struct sockaddr *)&sin, &addrlen);
if (ret) {
close(sock);
return LUA_HANDLE_ERROR(L, errno);
}
port = ntohs(sin.sin_port);
server_t *server = (server_t *)lua_newuserdata(L, sizeof(server_t));
memset(server, 0, sizeof(server_t));
server->sock = sock;
server->ip_address = sin.sin_addr.s_addr;
luaL_getmetatable(L, "ipc.server");
lua_setmetatable(L, -2);
lua_pushinteger(L, port);
return 2;
}
int flock_close(lua_State *L) {
int *handle = lua_touserdata(L, 1);
int fd = *handle;
if (fd) {
int ret = flock(fd, LOCK_UN);
close(fd);
*handle = 0;
if (ret < 0) return LUA_HANDLE_ERROR(L, errno);
}
return 0;
}
static int can_use_fastpath(lua_State *L, int sock, uint32_t bind_addr, uint32_t addr) {
#if defined(USE_CUDA) && !defined(__APPLE__)
if (bind_addr == addr) {
int device;
THCudaCheck(cudaGetDevice(&device));
int ret = send(sock, &device, sizeof(device), 0);
if (ret < 0) {
close(sock);
return LUA_HANDLE_ERROR(L, errno);
}
int remote_device;
ret = recv(sock, &remote_device, sizeof(remote_device), 0);
if (ret <= 0) {
close(sock);
return LUA_HANDLE_ERROR(L, errno);
}
if (device != remote_device) {
int can;
THCudaCheck(cudaDeviceCanAccessPeer(&can, device, remote_device));
if (can) {
cudaError_t err = cudaDeviceEnablePeerAccess(remote_device, 0);
if (err == cudaSuccess || err == cudaErrorPeerAccessAlreadyEnabled) {
if (err == cudaErrorPeerAccessAlreadyEnabled) cudaGetLastError();
fprintf(stderr, "INFO: torch-ipc: CUDA IPC enabled between GPU%d and GPU%d\n", device, remote_device);
return 1;
} else {
fprintf(stderr, "WARN: torch-ipc: CUDA IPC disabled between GPU%d and GPU%d: %s\n", device, remote_device, cudaGetErrorString(err));
}
} else {
fprintf(stderr, "INFO: torch-ipc: CUDA IPC not possible between GPU%d and GPU%d\n", device, remote_device);
}
}
}
#else
(void)L;
(void)sock;
(void)bind_addr;
(void)addr;
#endif
return 0;
}
static int destroy_server(lua_State *L, server_t *server) {
if (server->sock) {
int ret = close(server->sock);
if (ret) return LUA_HANDLE_ERROR(L, errno);
server->sock = 0;
}
client_t *client = server->clients;
while (client) {
client_t *next = client->next;
destroy_client(L, client);
client = next;
}
server->clients = NULL;
server->num_clients = 0;
destroy_copy_context(&server->copy_context);
return 0;
}
static int get_sockaddr(lua_State *L, const char *host, const char *port, struct sockaddr *addr, socklen_t *addrlen) {
struct addrinfo hints;
memset(&hints, 0, sizeof(hints));
hints.ai_family = AF_INET;
hints.ai_socktype = SOCK_STREAM;
hints.ai_protocol = IPPROTO_TCP;
struct addrinfo *ai;
int ret = getaddrinfo(host, port, &hints, &ai);
if (ret) return LUA_HANDLE_ERROR_STR(L, gai_strerror(ret));
if (*addrlen < ai->ai_addrlen) {
freeaddrinfo(ai);
return LUA_HANDLE_ERROR(L, ENOMEM);
}
memcpy(addr, ai->ai_addr, ai->ai_addrlen);
*addrlen = ai->ai_addrlen;
freeaddrinfo(ai);
return 0;
}
static int workqueue_queue_read(lua_State *L, queue_t *queue, int doNotBlock) {
pthread_mutex_lock(&queue->mutex);
while (1) {
if (queue->num_items) {
int ret = rb_load(L, queue->rb);
queue->num_items--;
pthread_cond_signal(&queue->write_avail_cond);
pthread_mutex_unlock(&queue->mutex);
if (ret < 0) return LUA_HANDLE_ERROR(L, ret);
return ret;
} else if (doNotBlock) {
break;
} else {
pthread_cond_wait(&queue->read_avail_cond, &queue->mutex);
}
}
pthread_mutex_unlock(&queue->mutex);
return 0;
}
int cliser_server_recv_any(lua_State *L) {
double t0 = _ipc_seconds();
server_t *server = (server_t *)lua_touserdata(L, 1);
const char *tag = luaL_optstring(L, 2, NULL);
fd_set fds;
FD_ZERO(&fds);
int highest = -1;
client_t *client = server->clients;
while (client) {
if (!tag || (client->tag && strcmp(tag, client->tag) == 0)) {
FD_SET(client->sock, &fds);
if (client->sock > highest) {
highest = client->sock;
}
}
client = client->next;
}
int ret = select(highest + 1, &fds, NULL, NULL, NULL);
if (ret < 0) return LUA_HANDLE_ERROR(L, errno);
client = server->clients;
while (client) {
if (!tag || (client->tag && strcmp(tag, client->tag) == 0)) {
if (FD_ISSET(client->sock, &fds)) {
ret = sock_recv_msg(L, client->sock, client->recv_rb, &server->copy_context);
if (ret == 1) {
server_client_t *server_client = (server_client_t *)lua_newuserdata(L, sizeof(server_client_t));
server_client->server = server;
server_client->client = client;
luaL_getmetatable(L, "ipc.server.client");
lua_setmetatable(L, -2);
ret = 2;
}
break;
}
}
client = client->next;
}
server->copy_context.rx.total_seconds += (_ipc_seconds() - t0);
server->copy_context.rx.num_calls++;
return ret;
}
static int destroy_client(lua_State *L, client_t *client) {
if (client->sock) {
int ret = close(client->sock);
if (ret) return LUA_HANDLE_ERROR(L, errno);
client->sock = 0;
}
if (client->send_rb) {
ringbuffer_destroy(client->send_rb);
client->send_rb = NULL;
}
if (client->recv_rb) {
ringbuffer_destroy(client->recv_rb);
client->recv_rb = NULL;
}
destroy_copy_context(&client->copy_context);
if (client->tag) {
free(client->tag);
client->tag = NULL;
}
free(client);
return 0;
}
static int sock_recv_raw(lua_State *L, int sock, void *ptr, size_t len, copy_context_t *copy_context) {
int ret = sock_recv(sock, ptr, len, copy_context);
if (ret < 0) return LUA_HANDLE_ERROR(L, errno);
if ((size_t)ret != len) return LUA_HANDLE_ERROR_STR(L, "failed to recv the correct number of bytes");
return 0;
}
int cliser_server_clients(lua_State *L) {
server_t *server = (server_t *)lua_touserdata(L, 1);
uint32_t wait;
const char *tag;
int fi;
int invert_order;
if (lua_type(L, 2) == LUA_TFUNCTION) {
wait = 0;
fi = 2;
if (lua_type(L, 3) == LUA_TSTRING) {
tag = luaL_optstring(L, 3, NULL);
invert_order = luaL_optinteger(L, 4, 0);
} else {
tag = NULL;
invert_order = luaL_optinteger(L, 3, 0);
}
} else {
wait = luaL_checkint(L, 2);
if (lua_type(L, 3) != LUA_TFUNCTION) return LUA_HANDLE_ERROR_STR(L, "expected a callback function as argument #3");
fi = 3;
tag = NULL;
invert_order = 0;
}
while (wait > server->num_clients) {
struct sockaddr addr;
socklen_t addrlen = sizeof(addr);
int ret = accept(server->sock, &addr, &addrlen);
if (ret <= 0) return LUA_HANDLE_ERROR(L, errno);
int sock = ret;
int value = 1;
ret = setsockopt(sock, IPPROTO_TCP, TCP_NODELAY, &value, sizeof(value));
if (ret) return LUA_HANDLE_ERROR(L, errno);
int use_fastpath = can_use_fastpath(L, sock, server->ip_address, ((struct sockaddr_in *)&addr)->sin_addr.s_addr);
client_t *client = (client_t *)calloc(1, sizeof(client_t));
client->sock = sock;
client->send_rb = ringbuffer_create(SEND_RECV_SIZE);
client->recv_rb = ringbuffer_create(SEND_RECV_SIZE);
client->copy_context.use_fastpath = use_fastpath;
insert_client(server, client);
}
client_t **clients = alloca(server->num_clients * sizeof(client_t*));
client_t *client = server->clients;
uint32_t i = 0;
while (client) {
if (!tag || (client->tag && strcmp(tag, client->tag) == 0)) {
clients[i] = client;
i++;
}
client = client->next;
}
if (invert_order) {
qsort(clients, i, sizeof(client_t*), compare_clients_inverted);
} else {
qsort(clients, i, sizeof(client_t*), compare_clients);
}
for (uint32_t j = 0; j < i; j++) {
lua_pushvalue(L, fi);
server_client_t *server_client = (server_client_t *)lua_newuserdata(L, sizeof(server_client_t));
server_client->server = server;
server_client->client = clients[j];
luaL_getmetatable(L, "ipc.server.client");
lua_setmetatable(L, -2);
lua_call(L, 1, 0);
server_client->server = NULL;
server_client->client = NULL;
}
lua_pushinteger(L, i);
return 1;
}
