Bugfix

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);

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
 }
 

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;
 }

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);

src/real.h

@@ -8,6 +8,8 @@
 
 #ifdef __linux__
 #include <sys/epoll.h>
+#include <sys/socket.h>
+#include <unistd.h>
 #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);

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;

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);