From f18db49de4d373e6b1eddafd555c621c6ebeb7fb Mon Sep 17 00:00:00 2001
From: GeorgyKirichenko <kirichenkoga@gmail.com>
Date: Wed, 18 Jan 2017 22:28:49 +0300
Subject: [PATCH] Add io wait callback for blocking io operations

Add a io wait callback that will be called instead of base io
writing handlers.
---
 include/ma_common.h        | 1 +
 include/mysql.h            | 3 ++-
 libmariadb/mariadb_lib.c   | 6 ++++++
 plugins/pvio/pvio_npipe.c  | 7 +++++++
 plugins/pvio/pvio_socket.c | 7 +++++++
 5 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/include/ma_common.h b/include/ma_common.h
index 01f7ba161..d5b268665 100644
--- a/include/ma_common.h
+++ b/include/ma_common.h
@@ -71,6 +71,7 @@ struct st_mysql_options_extension {
   char *server_public_key;
   char *proxy_header;
   size_t proxy_header_len;
+  int (*io_wait)(my_socket handle, my_bool is_read, int timeout);
 };
 
 typedef struct st_connection_handler
diff --git a/include/mysql.h b/include/mysql.h
index 08525e826..ef41db5f0 100644
--- a/include/mysql.h
+++ b/include/mysql.h
@@ -233,7 +233,8 @@ extern const char *SQLSTATE_UNKNOWN;
     MARIADB_OPT_MULTI_RESULTS,
     MARIADB_OPT_MULTI_STATEMENTS,
     MARIADB_OPT_INTERACTIVE,
-    MARIADB_OPT_PROXY_HEADER
+    MARIADB_OPT_PROXY_HEADER,
+    MARIADB_OPT_IO_WAIT
   };
 
   enum mariadb_value {
diff --git a/libmariadb/mariadb_lib.c b/libmariadb/mariadb_lib.c
index 3a4297878..495d1d0a2 100644
--- a/libmariadb/mariadb_lib.c
+++ b/libmariadb/mariadb_lib.c
@@ -3015,6 +3015,9 @@ mysql_optionsv(MYSQL *mysql,enum mysql_option option, ...)
   case MYSQL_OPT_TLS_VERSION:
     OPT_SET_EXTENDED_VALUE_STR(&mysql->options, tls_version, (char *)arg1);
     break;
+  case MARIADB_OPT_IO_WAIT:
+    mysql->options.extension->io_wait = (int(*)(my_socket, my_bool, int))arg1;
+    break;
   default:
     va_end(ap);
     return(-1);
@@ -3229,6 +3232,9 @@ mysql_get_optionv(MYSQL *mysql, enum mysql_option option, void *arg, ...)
   case MARIADB_OPT_CONNECTION_HANDLER:
     *((char **)arg)= mysql->options.extension ? mysql->options.extension->connection_handler : NULL;
     break;
+  case MARIADB_OPT_IO_WAIT:
+    *((int(**)(my_socket, my_bool, int))arg) = mysql->options.extension->io_wait;
+    break;
   default:
     va_end(ap);
     return(-1);
diff --git a/plugins/pvio/pvio_npipe.c b/plugins/pvio/pvio_npipe.c
index cf95ede8e..b9430310c 100644
--- a/plugins/pvio/pvio_npipe.c
+++ b/plugins/pvio/pvio_npipe.c
@@ -162,6 +162,13 @@ ssize_t pvio_npipe_write(MARIADB_PVIO *pvio, const uchar *buffer, size_t length)
 
 int pvio_npipe_wait_io_or_timeout(MARIADB_PVIO *pvio, my_bool is_read, int timeout)
 {
+  if (pvio->mysql->options.extension->io_wait != NULL) {
+    HANDLE handle;
+    if (pvio_npipe_get_handle(pvio, &handle))
+      return 0;
+    return pvio->mysql->options.extension->io_wait(handle, is_read, timeout);
+  }
+
   DWORD status;
   int save_error;
   struct st_pvio_npipe *cpipe= NULL;
diff --git a/plugins/pvio/pvio_socket.c b/plugins/pvio/pvio_socket.c
index 1936f3170..93251769e 100644
--- a/plugins/pvio/pvio_socket.c
+++ b/plugins/pvio/pvio_socket.c
@@ -504,6 +504,13 @@ int pvio_socket_wait_io_or_timeout(MARIADB_PVIO *pvio, my_bool is_read, int time
   if (!pvio || !pvio->data)
     return 0;
 
+  if (pvio->mysql->options.extension->io_wait != NULL) {
+    my_socket handle;
+    if (pvio_socket_get_handle(pvio, &handle))
+      return 0;
+    return pvio->mysql->options.extension->io_wait(handle, is_read, timeout);
+  }
+
   csock= (struct st_pvio_socket *)pvio->data;
   {
 #ifndef _WIN32