diff --git a/src/common.h b/src/common.h index 6ab3bb56..fbf39b60 100644 --- a/src/common.h +++ b/src/common.h @@ -127,6 +127,8 @@ static constexpr size_t kAltStackSize = 16 * 1024UL; // 16k sigaltstacks #define SIGQUIESCE (SIGRTMIN + 7) #define SIGDUMP (SIGRTMIN + 8) +static constexpr size_t kForkCopyFileSize = 32 * 1024 * 1024ul; + // BinnedTracker static constexpr size_t kBinnedTrackerBinCount = 1; static constexpr size_t kBinnedTrackerMaxEmpty = 128; @@ -235,6 +237,7 @@ inline mt19937_64 *initSeed() { unsigned long buf; auto sz = read(fd, (void *)&buf, sizeof(unsigned long)); hard_assert(sz == sizeof(unsigned long)); + close(fd); // std::random_device rd; // return new (mtBuf) std::mt19937_64(rd()); return new (mtBuf) std::mt19937_64(buf); diff --git a/src/libmesh.cc b/src/libmesh.cc index 73ee230f..06cc5461 100644 --- a/src/libmesh.cc +++ b/src/libmesh.cc @@ -227,6 +227,16 @@ int MESH_EXPORT epoll_pwait(int __epfd, struct epoll_event *__events, int __maxe return mesh::runtime().epollPwait(__epfd, __events, __maxevents, __timeout, __ss); } +#endif +#if __linux__ + +ssize_t MESH_EXPORT recv(int sockfd, void *buf, size_t len, int flags) { + return mesh::runtime().recv(sockfd, buf, len, flags); +} + +ssize_t MESH_EXPORT recvmsg(int sockfd, struct msghdr *msg, int flags) { + return mesh::runtime().recvmsg(sockfd, msg, flags); +} #endif } diff --git a/src/meshable_arena.cc b/src/meshable_arena.cc index 098749a1..7e2631c4 100644 --- a/src/meshable_arena.cc +++ b/src/meshable_arena.cc @@ -705,10 +705,13 @@ void MeshableArena::afterForkChild() { const int oldFd = _fd; - const auto bitmap = allocatedBitmap(); - for (auto const &i : bitmap) { - int result = internal::copyFile(newFd, oldFd, i * kPageSize, kPageSize); - d_assert(result == CPUInfo::PageSize); + const size_t end = _end * kPageSize; + for (size_t i = 0; i < end; i += kForkCopyFileSize) { + int result = internal::copyFile(newFd, oldFd, i, std::min(end - i, kForkCopyFileSize)); + d_assert(result >= 0); + } + + while (write(_forkPipe[1], "ok", strlen("ok")) == EAGAIN) { } int r = mprotect(_arenaBegin, kArenaSize, PROT_READ | PROT_WRITE); @@ -766,10 +769,7 @@ void MeshableArena::afterForkChild() { close(oldFd); - while (write(_forkPipe[1], "ok", strlen("ok")) == EAGAIN) { - } close(_forkPipe[1]); - _forkPipe[0] = -1; _forkPipe[1] = -1; } diff --git a/src/real.cc b/src/real.cc index d7f9dc69..763cda7f 100644 --- a/src/real.cc +++ b/src/real.cc @@ -21,6 +21,8 @@ namespace real { #ifdef __linux__ DEFINE_REAL(epoll_pwait); DEFINE_REAL(epoll_wait); +DEFINE_REAL(recv); +DEFINE_REAL(recvmsg); #endif DEFINE_REAL(pthread_create); @@ -40,6 +42,8 @@ void init() { #ifdef __linux__ INIT_REAL(epoll_pwait, RTLD_NEXT); INIT_REAL(epoll_wait, RTLD_NEXT); + INIT_REAL(recv, RTLD_NEXT); + INIT_REAL(recvmsg, RTLD_NEXT); #endif INIT_REAL(pthread_create, RTLD_NEXT); diff --git a/src/real.h b/src/real.h index faef57b3..5d2ad1e5 100644 --- a/src/real.h +++ b/src/real.h @@ -8,6 +8,8 @@ #ifdef __linux__ #include +#include +#include #endif #pragma once @@ -23,6 +25,8 @@ void init(); #ifdef __linux__ DECLARE_REAL(epoll_pwait); DECLARE_REAL(epoll_wait); +DECLARE_REAL(recv); +DECLARE_REAL(recvmsg); #endif DECLARE_REAL(pthread_create); diff --git a/src/runtime.cc b/src/runtime.cc index a0b82d0e..b44d8953 100644 --- a/src/runtime.cc +++ b/src/runtime.cc @@ -269,6 +269,35 @@ int Runtime::epollPwait(int __epfd, struct epoll_event *__events, int __maxevent return mesh::real::epoll_pwait(__epfd, __events, __maxevents, __timeout, __ss); } + +ssize_t Runtime::recv(int sockfd, void *buf, size_t len, int flags) { + if (unlikely(mesh::real::recv == nullptr)) { + mesh::real::init(); + } + ssize_t ret = mesh::real::recv(sockfd, buf, len, flags); + while (ret < 0 && errno == EFAULT && heap().okToProceed(buf)) { + ret = mesh::real::recv(sockfd, buf, len, flags); + } + return ret; +} + +ssize_t Runtime::recvmsg(int sockfd, struct msghdr *msg, int flags) { + if (unlikely(mesh::real::recvmsg == nullptr)) + mesh::real::init(); + + ssize_t ret = mesh::real::recvmsg(sockfd, msg, flags); + while (ret < 0 && errno == EFAULT && msg) { + for (size_t i = 0; i < msg->msg_iovlen; ++i) { + auto ptr = msg->msg_iov[i].iov_base; + if (ptr) { + heap().okToProceed(ptr); + } + } + ret = mesh::real::recvmsg(sockfd, msg, flags); + } + return ret; +} + #endif static struct sigaction sigbusAction; diff --git a/src/runtime.h b/src/runtime.h index f2ff4446..5cb13368 100644 --- a/src/runtime.h +++ b/src/runtime.h @@ -54,6 +54,8 @@ class Runtime { #ifdef __linux__ int epollWait(int __epfd, struct epoll_event *__events, int __maxevents, int __timeout); int epollPwait(int __epfd, struct epoll_event *__events, int __maxevents, int __timeout, const __sigset_t *__ss); + ssize_t recv(int sockfd, void *buf, size_t len, int flags); + ssize_t recvmsg(int sockfd, struct msghdr *msg, int flags); #endif int sigaction(int signum, const struct sigaction *act, struct sigaction *oldact);