Merge "shutdown socket gracefully with SocketDrainer" into idea133
diff --git a/Makefile.common b/Makefile.common
index e0123c8..b0f8594 100644
--- a/Makefile.common
+++ b/Makefile.common
@@ -121,6 +121,7 @@
 	android/base/containers/PodVector.cpp \
 	android/base/containers/StringVector.cpp \
 	android/base/files/PathUtils.cpp \
+	android/base/sockets/SocketDrainer.cpp \
 	android/base/Log.cpp \
 	android/base/String.cpp \
 	android/base/StringFormat.cpp \
diff --git a/android/base/sockets/SocketDrainer.cpp b/android/base/sockets/SocketDrainer.cpp
new file mode 100644
index 0000000..85dac36
--- /dev/null
+++ b/android/base/sockets/SocketDrainer.cpp
@@ -0,0 +1,226 @@
+// Copyright 2014 The Android Open Source Project
+//
+// This software is licensed under the terms of the GNU General Public
+// License version 2, as published by the Free Software Foundation, and
+// may be copied, distributed, and modified under those terms.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+
+#include <android/base/Limits.h>
+#include "android/base/sockets/SocketDrainer.h"
+
+#include "android/sockets.h"
+#include "android/utils/eintr_wrapper.h"
+#ifdef _WIN32
+#  include <winsock2.h>
+#else
+#  include <sys/socket.h>
+#endif
+
+#include <unistd.h>
+#include <set>
+
+// Some implementation whys:
+// When the looper is running, the sockets are non-blocking and are only
+// closed after all data has been read.
+// When looper stops, the remaining sockets are still non-blocking and are
+// closed after all available data has been read (cannot set to blocking, because
+// doing that could put the program sleeping forever).
+// When socket is already closed for read, we only need to shutdown writing
+// and close the socket.
+
+static android::base::SocketDrainer *s_socket_drainer = 0;
+
+namespace android {
+namespace base {
+
+class DrainerObject;
+
+// SocketDrainImpl implements the SocketDrainer and manages all the DrainerObjects
+class SocketDrainerImpl {
+public:
+    SocketDrainerImpl(Looper* looper) : mLooper(looper), mDrainerObjects() { }
+    ~SocketDrainerImpl() { drainAndCloseAllSockets(); }
+public:
+    void addSocketToDrain(int socket_fd);
+    void removeDrainerObject(DrainerObject* drainer) { mDrainerObjects.erase(drainer); }
+private:
+    Looper*  mLooper;
+    typedef std::set<DrainerObject*> DrainSet;
+    DrainSet mDrainerObjects;
+
+    void drainAndCloseAllSockets();
+};
+
+// DrainerObject drains and closes socket
+class DrainerObject {
+public:
+    DrainerObject(int socket_fd, Looper* looper, SocketDrainerImpl* parent);
+    ~DrainerObject();
+public:
+    // drain socket and return true if there is still more data to drain
+    bool drainSocket();
+
+    // return true when there is no data to drain
+    bool socketIsDrained() const { return mSocketIsDrained; }
+
+    // remove from SocketDrainerImpl
+    void removeFromParent() { if(mParent) mParent->removeDrainerObject(this); }
+private:
+    int     mSocket;
+    Looper* mLooper;
+    SocketDrainerImpl* mParent;
+    LoopIo  mIo[1];
+    bool    mSocketIsDrained;
+
+    void shutdownRead();
+    void shutdownWrite();
+    void closeSocket();
+};
+
+// callback from looper when the socket_fd has some data ready to read
+static void _on_read_socket_fd(void* opaque, int fd, unsigned events) {
+    DrainerObject * drainerObject = (DrainerObject*)opaque;
+    if (!drainerObject) return;
+    if ((events & LOOP_IO_READ) != 0) {
+        drainerObject->drainSocket();
+    }
+    if (drainerObject->socketIsDrained()) {
+        drainerObject->removeFromParent();
+        delete drainerObject;
+    }
+}
+
+DrainerObject::DrainerObject(int socket_fd, Looper* looper, SocketDrainerImpl* parent) :
+        mSocket(socket_fd), mLooper(looper), mParent(parent), mIo(), mSocketIsDrained(false) {
+    shutdownWrite();
+    if (drainSocket() && mLooper && mParent) {
+        loopIo_init(mIo, mLooper, mSocket, _on_read_socket_fd, this);
+        loopIo_wantRead(mIo);
+        loopIo_dontWantWrite(mIo);
+    } else {
+        // there is nothing to read, the drainer object is done
+        mLooper = 0;
+    }
+}
+
+DrainerObject::~DrainerObject() {
+    if (!mSocketIsDrained) {
+        char buff[1024];
+        while(socket_recv(mSocket, buff, sizeof(buff)) > 0) { ; }
+        mSocketIsDrained = true;
+    }
+    shutdownRead();
+    if (mLooper) {
+        loopIo_dontWantRead(mIo);
+        loopIo_done(mIo);
+        mLooper = 0;
+    }
+    closeSocket();
+    mSocket = -1;
+    mParent = 0;
+}
+
+bool DrainerObject::drainSocket() {
+    errno = 0;
+    char buff[1024];
+    int size = socket_recv(mSocket, buff, sizeof(buff));
+    if (size > 0) {
+        return true;
+    } else if (size < 0 && errno == EWOULDBLOCK) {
+        return true;
+    }
+    mSocketIsDrained = true;
+    return false;
+}
+
+void DrainerObject::shutdownWrite() {
+#ifdef _WIN32
+    shutdown(mSocket, SD_SEND);
+#else
+    shutdown(mSocket, SHUT_WR);
+#endif
+}
+
+void DrainerObject::shutdownRead() {
+#ifdef _WIN32
+    shutdown(mSocket, SD_RECEIVE);
+#else
+    shutdown(mSocket, SHUT_RD);
+#endif
+}
+
+void DrainerObject::closeSocket() {
+#ifdef _WIN32
+    closesocket(mSocket);
+#else
+    IGNORE_EINTR(close(mSocket));
+#endif
+}
+
+//--------------------------- SocketDrainerImpl Implementation -------------------------
+
+void SocketDrainerImpl::addSocketToDrain(int socket_fd) {
+    DrainerObject* drainer = new DrainerObject(socket_fd, mLooper, this);
+    if (drainer->socketIsDrained()) {
+        delete drainer;
+    } else {
+        mDrainerObjects.insert(drainer);
+    }
+}
+
+void SocketDrainerImpl::drainAndCloseAllSockets() {
+    for (DrainSet::iterator it = mDrainerObjects.begin(); it != mDrainerObjects.end(); ++it) {
+        DrainerObject* drainer = *it;
+        delete drainer;
+    }
+    mDrainerObjects.clear();
+}
+
+//--------------------------- SocketDrainer Implementation -----------------------------
+
+SocketDrainer::SocketDrainer(Looper* looper) :
+    mSocketDrainerImpl(new SocketDrainerImpl(looper)) {
+}
+
+SocketDrainer::~SocketDrainer() {
+    delete mSocketDrainerImpl;
+    mSocketDrainerImpl = 0;
+}
+
+void
+SocketDrainer::drainAndClose(int socket_fd) {
+    if (socket_fd < 0) return;
+    mSocketDrainerImpl->addSocketToDrain(socket_fd);
+}
+
+} // namespace base
+} // namespace android
+
+// -------------------- extern C functions ---------------------------------------------
+
+void socket_drainer_start(Looper* looper) {
+    if (!looper) return;
+    if (!s_socket_drainer) s_socket_drainer = new android::base::SocketDrainer(looper);
+}
+
+void socket_drainer_drain_and_close(int socket_fd) {
+    if (socket_fd < 0) return;
+    socket_set_nonblock(socket_fd);
+    if (s_socket_drainer) {
+        s_socket_drainer->drainAndClose(socket_fd);
+    } else {
+        android::base::DrainerObject drainerObject(socket_fd, 0, 0);
+    }
+}
+
+void socket_drainer_stop() {
+    if (s_socket_drainer) {
+        delete s_socket_drainer;
+        s_socket_drainer = 0;
+    }
+}
diff --git a/android/base/sockets/SocketDrainer.h b/android/base/sockets/SocketDrainer.h
new file mode 100644
index 0000000..a9f6b7f
--- /dev/null
+++ b/android/base/sockets/SocketDrainer.h
@@ -0,0 +1,73 @@
+// Copyright 2014 The Android Open Source Project
+//
+// This software is licensed under the terms of the GNU General Public
+// License version 2, as published by the Free Software Foundation, and
+// may be copied, distributed, and modified under those terms.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+#ifndef ANDROID_BASE_SOCKETS_SOCKET_DRAINER_H
+#define ANDROID_BASE_SOCKETS_SOCKET_DRAINER_H
+
+#include <android/utils/compiler.h>
+
+/*
+This class is used to manage graceful closure of sockets, as described by
+"The ultimate SO_LINGER page", which requires several steps to do properly.
+
+To use this class from C code, call the following after creating a looper:
+   socket_drainer_start(looper);
+
+When a socket needs to be gracefully closed, call the following:
+   socket_drainer_drain_and_close(socket_fd);
+
+After looper_run completes, call the following:
+   socket_drainer_stop();
+
+Note: the "socket_drainer_drain_and_close" function can be called before
+"socket_drainer_start" or after "socket_drainer_stop" . In those cases,
+the socket is simply shut down and closed without receiving all the pending data.
+
+For C++ code, simply create an object as follows:
+   SocketDrainer socketDrainer(looper);
+
+When a socket needs to be gracefully closed, call the following:
+   socketDrainer.drainAndClose(socket_fd);
+
+When the socketDrainer is destroyed (for example, falling out of scope), all remaining
+sockets will be shut down and closed without futher draining.
+
+*/
+
+ANDROID_BEGIN_HEADER
+
+#include "android/looper.h"
+
+void socket_drainer_start(Looper* looper);
+void socket_drainer_drain_and_close(int socket_fd);
+void socket_drainer_stop();
+
+ANDROID_END_HEADER
+
+#ifdef __cplusplus
+namespace android {
+namespace base {
+
+class SocketDrainerImpl;
+class SocketDrainer {
+public:
+    SocketDrainer(Looper* looper);
+    ~SocketDrainer();
+    // Add a socket to be drained and closed
+    void drainAndClose(int socket_fd);
+private:
+    SocketDrainerImpl*  mSocketDrainerImpl;
+};
+
+} // namespace base
+} // namespace android
+#endif
+#endif
diff --git a/android/sockets.c b/android/sockets.c
index 4b8796d..c3da5f1 100644
--- a/android/sockets.c
+++ b/android/sockets.c
@@ -21,6 +21,7 @@
 #include "sysemu/char.h"
 #include <stdlib.h>
 #include <string.h>
+#include "android/base/sockets/SocketDrainer.h"
 #include "android/utils/path.h"
 #include "android/utils/debug.h"
 #include "android/utils/eintr_wrapper.h"
@@ -1230,57 +1231,12 @@
 
 #endif /* !_WIN32 */
 
-#ifdef _WIN32
-
-static void
-socket_close_handler( void*  _fd )
-{
-    int   fd = (int)_fd;
-    int   ret;
-    char  buff[64];
-
-    /* we want to drain the read side of the socket before closing it */
-    do {
-        ret = recv( fd, buff, sizeof(buff), 0 );
-    } while (ret < 0 && WSAGetLastError() == WSAEINTR);
-
-    if (ret < 0 && WSAGetLastError() == EWOULDBLOCK)
-        return;
-
-    qemu_set_fd_handler( fd, NULL, NULL, NULL );
-    closesocket( fd );
-}
 
 void
-socket_close( int  fd )
-{
-    int  old_errno = errno;
-
-    shutdown( fd, SD_BOTH );
-    /* we want to drain the socket before closing it */
-    qemu_set_fd_handler( fd, socket_close_handler, NULL, (void*)fd );
-
-    errno = old_errno;
+socket_close( int  fd ) {
+    socket_drainer_drain_and_close(fd);
 }
 
-#else /* !_WIN32 */
-
-#include <unistd.h>
-
-void
-socket_close( int  fd )
-{
-    int  old_errno = errno;
-
-    shutdown( fd, SHUT_RDWR );
-    IGNORE_EINTR(close( fd ));
-
-    errno = old_errno;
-}
-
-#endif /* !_WIN32 */
-
-
 static int
 socket_bind_server( int  s, const SockAddress*  to, SocketType  type )
 {
diff --git a/include/android/sockets.h b/include/android/sockets.h
index 47d6c91..b15922f 100644
--- a/include/android/sockets.h
+++ b/include/android/sockets.h
@@ -13,10 +13,13 @@
 #ifndef QEMU_SOCKET_H
 #define QEMU_SOCKET_H
 
+#include <android/utils/compiler.h>
 #include <stddef.h>
 #include <stdint.h>
 #include <errno.h>
 
+ANDROID_BEGIN_HEADER
+
 /* we're going to hide the implementation details of sockets behind
  * a simple wrapper interface declared here.
  *
@@ -420,4 +423,6 @@
 int  socket_mcast_inet_set_loop( int  s, int  enabled );
 int  socket_mcast_inet_set_ttl( int  s, int  ttl );
 
+ANDROID_END_HEADER
+
 #endif /* QEMU_SOCKET_H */
diff --git a/vl-android.c b/vl-android.c
index be33559..a1ac5c1 100644
--- a/vl-android.c
+++ b/vl-android.c
@@ -50,6 +50,7 @@
 
 #include "migration/qemu-file.h"
 #include "android/android.h"
+#include "android/base/sockets/SocketDrainer.h"
 #include "android/charpipe.h"
 #include "android/log-rotate.h"
 #include "modem_driver.h"
@@ -2382,6 +2383,8 @@
     android_hw_control_init();
     android_net_pipes_init();
 
+    socket_drainer_start(looper_newCore());
+
 #ifdef CONFIG_KVM
     /* By default, force auto-detection for kvm */
     kvm_allowed = -1;
@@ -4235,6 +4238,8 @@
     main_loop();
     quit_timers();
     net_cleanup();
+    socket_drainer_stop();
+
     android_emulation_teardown();
     return 0;
 }