perf: Copy files instead of sending over HTTP if Wave app and Wave server are running on the same machine. #982

file_server.go

@@ -19,6 +19,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"mime"
 	"mime/multipart"
 	"net/http"
 	"os"
@@ -182,7 +183,15 @@ func (fs *FileServer) storeFilesInSingleDir(files []*multipart.FileHeader) ([]st
 		}
 		defer src.Close()
 
-		dir, file := filepath.Split(file.Filename)
+		// Need to parse the filename from the Content-Disposition header due to HTTP standard saying FileName should be basename.
+		// https://github.com/golang/go/blob/8dbf3e9393400d72d313e5616c88873e07692c70/src/mime/multipart/multipart.go#L82-L84
+		_, params, _ := mime.ParseMediaType(file.Header.Get("Content-Disposition"))
+		filename := params["filename"]
+		if filename == "" {
+			filename = file.Filename
+		}
+
+		dir, file := filepath.Split(filename)
 		uploadPath := filepath.Join(uploadDir, dir)
 
 		if err := os.MkdirAll(uploadPath, 0700); err != nil {

py/Makefile

@@ -37,6 +37,8 @@ docs: ## Build API docs
 test:
 	./venv/bin/python -m tests
 	echo "Testing using BASE_URL" && H2O_WAVE_BASE_URL="/foo/" ./venv/bin/python -m tests
+	echo "Testing using LOCAL UPLOAD" && H2O_WAVE_WAVED_DIR=".." ./venv/bin/python -m tests
+	echo "Testing using LOCAL UPLOAD AND BASE URL" && H2O_WAVE_WAVED_DIR=".." H2O_WAVE_BASE_URL="/foo/" ./venv/bin/python -m tests
 
 purge: ## Purge previous build
 	rm -rf build dist h2o_wave.egg-info

py/h2o_wave/cli.py

@@ -137,6 +137,8 @@ def run(app: str, no_reload: bool, no_autostart: bool):
     finally:
         if not server_not_running:
             try:
+                if not os.environ.get('H2O_WAVE_WAVED_DIR') and is_waved_present:
+                    os.environ['H2O_WAVE_WAVED_DIR'] = sys.exec_prefix
                 uvicorn.run(f'{app}:main', host=_localhost, port=port, reload=not no_reload)
             except Exception as e:
                 if waved_process:

py/h2o_wave/core.py

@@ -13,8 +13,15 @@
 # limitations under the License.
 
 from io import BufferedReader
+import asyncio
+import ipaddress
 import json
+import platform
 import secrets
+import shutil
+import subprocess
+from urllib.parse import urlparse
+from uuid import uuid4
 import warnings
 import logging
 import os
@@ -40,6 +47,7 @@ def _get_env(key: str, value: Any):
 
 
 _default_internal_address = 'http://127.0.0.1:8000'
+_base_url = _get_env('BASE_URL', '/')
 
 
 class _Config:
@@ -657,17 +665,47 @@ def upload(self, files: List[str]) -> List[str]:
         Returns:
             A list of remote URLs for the uploaded files, in order.
         """
-        upload_files = []
-        file_handles: List[BufferedReader] = []
         for f in files:
-            file_handle = open(f, 'rb')
-            upload_files.append(('files', (os.path.basename(f), file_handle)))
-            file_handles.append(file_handle)
+            if not os.path.isfile(f):
+                raise ValueError(f'{f} is not a file.')
+
+        waved_dir = _get_env('WAVED_DIR', None)
+        data_dir = _get_env('DATA_DIR', 'data')
+        skip_local_upload = _get_env('NO_COPY_UPLOAD', 'false').lower() in ['true', '1', 't']
+
+        # If we know the path of waved and running app on the same machine,
+        # we can simply copy the files instead of making an HTTP request.
+        if _is_loopback_address() and not skip_local_upload and waved_dir and data_dir:
+            try:
+                uploaded_files = []
+                for f in files:
+                    uuid = str(uuid4())
+                    dst = os.path.join(waved_dir, data_dir, 'f', uuid)
+                    os.makedirs(dst, exist_ok=True)
+
+                    if 'Windows' in platform.system():
+                        src = os.path.dirname(f) or os.getcwd()
+                        args = ['robocopy', src, dst, os.path.basename(f), '/J', '/W:0']
+                    else:
+                        args = ['cp', f, dst]
+
+                    _, err = subprocess.Popen(args, stderr=subprocess.PIPE, stdout=subprocess.DEVNULL).communicate()
+                    if err:
+                        raise ValueError(err.decode())
+
+                    uploaded_files.append(f'{_base_url}_f/{uuid}/{os.path.basename(f)}')
+                return uploaded_files
+            except:
+                pass
+
+        uploaded_files = []
+        for f in files:
+            uploaded_files.append(('files', (os.path.basename(f), open(f, 'rb'))))
 
-        res = self._http.post(f'{_config.hub_address}_f/', files=upload_files)
+        res = self._http.post(f'{_config.hub_address}_f/', files=uploaded_files)
 
-        for h in file_handles:
-            h.close()
+        for _, f in uploaded_files:
+            f[1].close()
 
         if res.status_code == 200:
             return json.loads(res.text)['files']
@@ -687,17 +725,39 @@ def upload_dir(self, directory: str) -> str:
         if not os.path.isdir(directory):
             raise ValueError(f'{directory} is not a directory.')
 
+        waved_dir = _get_env('WAVED_DIR', None)
+        data_dir = _get_env('DATA_DIR', 'data')
+        skip_local_upload = _get_env('NO_COPY_UPLOAD', 'false').lower() in ['true', '1', 't']
+
+        # If we know the path of waved and running app on the same machine,
+        # we can simply copy the files instead of making an HTTP request.
+        if _is_loopback_address() and not skip_local_upload and waved_dir and data_dir:
+            try:
+                uuid = str(uuid4())
+                dst = os.path.join(waved_dir, data_dir, 'f', uuid)
+                os.makedirs(dst, exist_ok=True)
+
+                if 'Windows' in platform.system():
+                    args = ['robocopy', directory, dst, '/S', '/J', '/W:0', '*.*']
+                else:
+                    args = ['rsync', '-a', os.path.join(directory, '.'), dst]
+
+                _, err = subprocess.Popen(args, stderr=subprocess.PIPE, stdout=subprocess.DEVNULL).communicate()
+                if err:
+                    raise ValueError(err.decode())
+
+                return [f'{_base_url}_f/{uuid}']
+            except:
+                pass
+
         upload_files = []
-        file_handles: List[BufferedReader] = []
         for f in _get_files_in_directory(directory, []):
-            file_handle = open(f, 'rb')
-            upload_files.append(('files', (os.path.relpath(f, directory), file_handle)))
-            file_handles.append(file_handle)
+            upload_files.append(('files', (os.path.relpath(f, directory), open(f, 'rb'))))
 
         res = self._http.post(f'{_config.hub_address}_f/', headers={'Wave-Directory-Upload': "True"}, files=upload_files)
 
-        for h in file_handles:
-            h.close()
+        for _, f in upload_files:
+            f[1].close()
 
         if res.status_code == 200:
             return json.loads(res.text)['files']
@@ -841,17 +901,35 @@ async def upload_dir(self, directory: str) -> str:
         if not os.path.isdir(directory):
             raise ValueError(f'{directory} is not a directory.')
 
+        waved_dir = _get_env('WAVED_DIR', None)
+        data_dir = _get_env('DATA_DIR', 'data')
+        skip_local_upload = _get_env('NO_COPY_UPLOAD', 'false').lower() in ['true', '1', 't']
+
+        # If we know the path of waved and running app on the same machine,
+        # we can simply copy the files instead of making an HTTP request.
+        if _is_loopback_address() and not skip_local_upload and waved_dir and data_dir:
+            try:
+                uuid = str(uuid4())
+                dst = os.path.join(waved_dir, data_dir, 'f', uuid)
+                os.makedirs(dst, exist_ok=True)
+
+                if 'Windows' in platform.system():
+                    args = ['robocopy', directory, dst, '/S', '/J', '/W:0', '*.*']
+                else:
+                    args = ['rsync', '-a', os.path.join(directory, '.'), dst]
+
+                return [await _copy_in_subprocess(args, uuid)]
+            except:
+                pass
+
         upload_files = []
-        file_handles: List[BufferedReader] = []
         for f in _get_files_in_directory(directory, []):
-            file_handle = open(f, 'rb')
-            upload_files.append(('files', (os.path.relpath(f, directory), file_handle)))
-            file_handles.append(file_handle)
+            upload_files.append(('files', (os.path.relpath(f, directory), open(f, 'rb'))))
 
         res = await self._http.post(f'{_config.hub_address}_f/', headers={'Wave-Directory-Upload': "True"}, files=upload_files)
 
-        for h in file_handles:
-            h.close()
+        for _, f in upload_files:
+            f[1].close()
 
         if res.status_code == 200:
             return json.loads(res.text)['files']
@@ -867,6 +945,36 @@ async def upload(self, files: List[str]) -> List[str]:
         Returns:
             A list of remote URLs for the uploaded files, in order.
         """
+        for f in files:
+            if not os.path.isfile(f):
+                raise ValueError(f'{f} is not a file.')
+
+        waved_dir = _get_env('WAVED_DIR', None)
+        data_dir = _get_env('DATA_DIR', 'data')
+        skip_local_upload = _get_env('NO_COPY_UPLOAD', 'false').lower() in ['true', '1', 't']
+
+        # If we know the path of waved and running app on the same machine,
+        # we can simply copy the files instead of making an HTTP request.
+        if _is_loopback_address() and not skip_local_upload and waved_dir and data_dir:
+            try:
+                tasks = []
+                for f in files:
+                    uuid = str(uuid4())
+                    dst = os.path.join(waved_dir, data_dir, 'f', uuid)
+                    os.makedirs(dst, exist_ok=True)
+
+                    if 'Windows' in platform.system():
+                        src = os.path.dirname(f) or os.getcwd()
+                        args = ['robocopy', src, dst, os.path.basename(f), '/J', '/W:0']
+                    else:
+                        args = ['cp', f, dst]
+
+                    tasks.append(asyncio.create_task(_copy_in_subprocess(args, uuid, f)))
+
+                return await asyncio.gather(*tasks)
+            except:
+                pass
+
         upload_files = []
         file_handles: List[BufferedReader] = []
         for f in files:
@@ -896,6 +1004,7 @@ async def download(self, url: str, path: str) -> str:
         path = os.path.abspath(path)
         # If path is a directory, get basename from url
         filepath = os.path.join(path, os.path.basename(url)) if os.path.isdir(path) else path
+
         async with self._http.stream('GET', f'{_config.hub_host_address}{url}') as res:
             if res.status_code != 200:
                 await res.aread()
@@ -963,6 +1072,19 @@ async def proxy(self, method: str, url: str, headers: Optional[Dict[str, List[st
         raise ServiceError(f'Proxy request failed (code={res.status_code}): {res.text}')
 
 
+async def _copy_in_subprocess(args: List[str], uuid: str, f='') -> str:
+    p = await asyncio.create_subprocess_exec(*args, stderr=subprocess.PIPE, stdout=subprocess.DEVNULL)
+    _, err = await p.communicate()
+
+    if err:
+        raise ValueError(err.decode())
+
+    if f:
+        return f'{_base_url}_f/{uuid}/{os.path.basename(f)}'
+    else:
+        return f'{_base_url}_f/{uuid}'
+
+
 def _get_files_in_directory(directory: str, files: List[str]) -> List[str]:
     for f in os.listdir(directory):
         path = os.path.join(directory, f)
@@ -972,6 +1094,7 @@ def _get_files_in_directory(directory: str, files: List[str]) -> List[str]:
             _get_files_in_directory(path, files)
     return files
 
+
 def marshal(d: Any) -> str:
     """
     Marshal to JSON.
@@ -1008,3 +1131,11 @@ def pack(data: Any) -> str:
     The object or value compressed into a string.
     """
     return 'data:' + marshal(_dump(data))
+
+
+def _is_loopback_address() -> bool:
+    try:
+        hostname = urlparse(_config.hub_address).hostname
+        return ipaddress.ip_address(hostname).is_loopback
+    except ValueError:
+        return False

py/setup.py

@@ -36,6 +36,7 @@
 __arch__ = "{arch}"
 ''')
 
+
 def get_data_files():
     data_dict = dict()
     data_dict['project_templates'] = [os.path.join('project_templates', f) for f in os.listdir('project_templates')]

py/tests/test_python_server.py

@@ -330,7 +330,6 @@ def test_cyc_buf_write(self):
             i=2,
         ))))
 
-
     def test_proxy(self):
         # waved -proxy must be set
         url = 'https://wave.h2o.ai'
@@ -342,7 +341,6 @@ def test_proxy(self):
             assert result.code == 400
             assert len(result.headers) > 0
 
-
     def test_file_server(self):
         f1 = 'temp_file1.txt'
         with open(f1, 'w') as f:
@@ -356,14 +354,12 @@ def test_file_server(self):
         os.remove(f2)
         assert s1 == s2
 
-
     def test_public_dir(self):
         p = site.download(f'{base_url}assets/brand/h2o.svg', 'h2o.svg')
         svg = read_file(p)
         os.remove(p)
         assert svg.index('<svg') == 0
 
-
     def test_cache(self):
         d1 = dict(foo='bar', qux=42)
         site.cache.set('test', 'data', d1)
@@ -375,7 +371,6 @@ def test_cache(self):
         assert d2['foo'] == d1['foo']
         assert d2['qux'] == d1['qux']
 
-
     def test_multipart_server(self):
         file_handle = open('../assets/brand/wave.svg', 'rb')
         p = site.uplink('test_stream', 'image/svg+xml', file_handle)
@@ -385,14 +380,13 @@ def test_multipart_server(self):
 
     def test_upload_dir(self):
         upload_path, = site.upload_dir(os.path.join('tests', 'test_folder'))
-        download_path = site.download(f'{upload_path}/test.txt', 'test.txt')
+        download_path = site.download(f'{upload_path}/dir1/test.txt', 'test.txt')
         txt = read_file(download_path)
         os.remove(download_path)
         assert len(txt) > 0
 
-
     def test_deleting_files(self):
-        upload_path, = site.upload([os.path.join('tests', 'test_folder', 'test.txt')])
+        upload_path, = site.upload([os.path.join('tests', 'test_folder', 'dir1', 'test.txt')])
         res = httpx.get(f'http://localhost:10101{upload_path}')
         assert res.status_code == 200
         site.unload(upload_path)