#830 Implement a socket wrapper API for use with LSP connections

7 files changed, 340 insertions, 36 deletions

diff --git a/autoload/ale/job.vim b/autoload/ale/job.vim
index 6ffc2a06..e0266cba 100644
--- a/autoload/ale/job.vim
+++ b/autoload/ale/job.vim
@@ -26,34 +26,11 @@ function! s:KillHandler(timer) abort
     call job_stop(l:job, 'kill')
 endfunction
 
-" Note that jobs and IDs are the same thing on NeoVim.
-function! ale#job#JoinNeovimOutput(job, last_line, data, mode, callback) abort
-    if a:mode is# 'raw'
-        call a:callback(a:job, join(a:data, "\n"))
-        return ''
-    endif
-
-    let l:lines = a:data[:-2]
-
-    if len(a:data) > 1
-        let l:lines[0] = a:last_line . l:lines[0]
-        let l:new_last_line = a:data[-1]
-    else
-        let l:new_last_line = a:last_line . get(a:data, 0, '')
-    endif
-
-    for l:line in l:lines
-        call a:callback(a:job, l:line)
-    endfor
-
-    return l:new_last_line
-endfunction
-
 function! s:NeoVimCallback(job, data, event) abort
     let l:info = s:job_map[a:job]
 
     if a:event is# 'stdout'
-        let l:info.out_cb_line = ale#job#JoinNeovimOutput(
+        let l:info.out_cb_line = ale#util#JoinNeovimOutput(
         \   a:job,
         \   l:info.out_cb_line,
         \   a:data,
@@ -61,7 +38,7 @@ function! s:NeoVimCallback(job, data, event) abort
         \   ale#util#GetFunction(l:info.out_cb),
         \)
     elseif a:event is# 'stderr'
-        let l:info.err_cb_line = ale#job#JoinNeovimOutput(
+        let l:info.err_cb_line = ale#util#JoinNeovimOutput(
         \   a:job,
         \   l:info.err_cb_line,
         \   a:data,
diff --git a/autoload/ale/socket.vim b/autoload/ale/socket.vim
new file mode 100644
index 00000000..78eba737
--- /dev/null
+++ b/autoload/ale/socket.vim
@@ -0,0 +1,137 @@
+" Author: w0rp <devw0rp@gmail.com>
+" Description: APIs for working with asynchronous sockets, with an API
+" normalised between Vim 8 and NeoVim. Socket connections only work in NeoVim
+" 0.3+, and silently do nothing in earlier NeoVim versions.
+"
+" Important functions are described below. They are:
+"
+"   ale#socket#Open(address, options) -> channel_id (>= 0 if successful)
+"   ale#socket#IsOpen(channel_id) -> 1 if open, 0 otherwise
+"   ale#socket#Close(channel_id)
+"   ale#socket#Send(channel_id, data)
+
+let s:channel_map = get(s:, 'channel_map', {})
+
+function! s:VimOutputCallback(channel, data) abort
+    let l:channel_id = ch_info(a:channel).id
+
+    " Only call the callbacks for jobs which are valid.
+    if l:channel_id >= 0 && has_key(s:channel_map, l:channel_id)
+        call ale#util#GetFunction(s:channel_map[l:channel_id].callback)(l:channel_id, a:data)
+    endif
+endfunction
+
+function! s:NeoVimOutputCallback(channel_id, data, event) abort
+    let l:info = s:channel_map[a:channel_id]
+
+    if a:event is# 'data'
+        let l:info.last_line = ale#util#JoinNeovimOutput(
+        \   a:channel_id,
+        \   l:info.last_line,
+        \   a:data,
+        \   l:info.mode,
+        \   ale#util#GetFunction(l:info.callback),
+        \)
+    endif
+endfunction
+
+" Open a socket for a given address. The following options are accepted:
+"
+" callback - A callback for receiving input. (required)
+"
+" A non-negative number representing a channel ID will be returned is the
+" connection was successful. 0 is a valid channel ID in Vim, so test if the
+" connection ID is >= 0.
+function! ale#socket#Open(address, options) abort
+    let l:mode = get(a:options, 'mode', 'raw')
+    let l:Callback = a:options.callback
+
+    let l:channel_info = {
+    \   'mode': l:mode,
+    \   'callback': a:options.callback,
+    \}
+
+    if !has('nvim')
+        " Vim
+        let l:channel_info.channel = ch_open(a:address, {
+        \   'mode': l:mode,
+        \   'waittime': 0,
+        \   'callback': function('s:VimOutputCallback'),
+        \})
+        let l:vim_info = ch_info(l:channel_info.channel)
+        let l:channel_id = !empty(l:vim_info) ? l:vim_info.id : -1
+    elseif exists('*chansend') && exists('*sockconnect')
+        " NeoVim 0.3+
+        try
+            let l:channel_id = sockconnect('tcp', a:address, {
+            \   'on_data': function('s:NeoVimOutputCallback'),
+            \})
+            let l:channel_info.last_line = ''
+        catch /connection failed/
+            let l:channel_id = -1
+        endtry
+
+        " 0 means the connection failed some times in NeoVim, so make the ID
+        " invalid to match Vim.
+        if l:channel_id is 0
+            let l:channel_id = -1
+        endif
+
+        let l:channel_info.channel = l:channel_id
+    else
+        " Other Vim versions.
+        let l:channel_id = -1
+    endif
+
+    if l:channel_id >= 0
+        let s:channel_map[l:channel_id] = l:channel_info
+    endif
+
+    return l:channel_id
+endfunction
+
+" Return 1 is a channel is open, 0 otherwise.
+function! ale#socket#IsOpen(channel_id) abort
+    if !has_key(s:channel_map, a:channel_id)
+        return 0
+    endif
+
+    if has('nvim')
+        " In NeoVim, we have to check if this channel is in the global list.
+        return index(map(nvim_list_chans(), 'v:val.id'), a:channel_id) >= 0
+    endif
+
+    let l:channel = s:channel_map[a:channel_id].channel
+    return ch_status(l:channel) is# 'open'
+endfunction
+
+" Close a socket, if it's still open.
+function! ale#socket#Close(channel_id) abort
+    " IsRunning isn't called here, so we don't check nvim_list_chans()
+    if !has_key(s:channel_map, a:channel_id)
+        return 0
+    endif
+
+    let l:channel = remove(s:channel_map, a:channel_id).channel
+
+    if has('nvim')
+        silent! call chanclose(l:channel)
+    elseif ch_status(l:channel) is# 'open'
+        call ch_close(l:channel)
+    endif
+endfunction
+
+" Send some data to a socket.
+function! ale#socket#Send(channel_id, data) abort
+    if !has_key(s:channel_map, a:channel_id)
+        return
+    endif
+
+    let l:channel = s:channel_map[a:channel_id].channel
+
+    if has('nvim')
+        call chansend(l:channel, a:data)
+    else
+        call ch_sendraw(l:channel, a:data)
+    endif
+endfunction
diff --git a/autoload/ale/util.vim b/autoload/ale/util.vim
index 28ab8231..d7b6904c 100644
--- a/autoload/ale/util.vim
+++ b/autoload/ale/util.vim
@@ -46,6 +46,33 @@ if !exists('g:ale#util#nul_file')
     endif
 endif
 
+" Given a job, a buffered line of data, a list of parts of lines, a mode data
+" is being read in, and a callback, join the lines of output for a NeoVim job
+" or socket together, and call the callback with the joined output.
+"
+" Note that jobs and IDs are the same thing on NeoVim.
+function! ale#util#JoinNeovimOutput(job, last_line, data, mode, callback) abort
+    if a:mode is# 'raw'
+        call a:callback(a:job, join(a:data, "\n"))
+        return ''
+    endif
+
+    let l:lines = a:data[:-2]
+
+    if len(a:data) > 1
+        let l:lines[0] = a:last_line . l:lines[0]
+        let l:new_last_line = a:data[-1]
+    else
+        let l:new_last_line = a:last_line . get(a:data, 0, '')
+    endif
+
+    for l:line in l:lines
+        call a:callback(a:job, l:line)
+    endfor
+
+    return l:new_last_line
+endfunction
+
 " Return the number of lines for a given buffer.
 function! ale#util#GetLineCount(buffer) abort
     return len(getbufline(a:buffer, 1, '$'))
diff --git a/test/dumb_tcp_client.py b/test/dumb_tcp_client.py
new file mode 100644
index 00000000..3a728b02
--- /dev/null
+++ b/test/dumb_tcp_client.py
@@ -0,0 +1,33 @@
+"""
+This is just a script for testing that the dumb TCP server actually works
+correctly, for verifying that problems with tests are in Vim. Pass the
+same port number given to the test server to check that it's working.
+"""
+import socket
+import sys
+
+
+def main():
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    result = sock.connect_ex(('127.0.0.1', int(sys.argv[1])))
+
+    if result:
+        sock.close()
+        sys.exit("Couldn't connect to the socket!")
+
+    data_sent = 'x' * 1024
+
+    sock.send(data_sent)
+    data_received = sock.recv(1024)
+
+    if data_sent != data_received:
+        sock.close()
+        sys.exit("Data sent didn't match data received.")
+
+    sock.close()
+
+    print("Everything was just fine.")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/test/dumb_tcp_server.py b/test/dumb_tcp_server.py
new file mode 100644
index 00000000..c15db65e
--- /dev/null
+++ b/test/dumb_tcp_server.py
@@ -0,0 +1,40 @@
+"""
+This Python script creates a TCP server that does nothing but send its input
+back to the client that connects to it. Only one argument must be given, a port
+to bind to.
+"""
+import os
+import socket
+import sys
+
+
+def main():
+    if len(sys.argv) < 2 or not sys.argv[1].isdigit():
+        sys.exit('You must specify a port number')
+
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    sock.bind(('127.0.0.1', int(sys.argv[1])))
+    sock.listen(0)
+
+    pid = os.fork()
+
+    if pid:
+        print(pid)
+        sys.exit()
+
+    while True:
+        connection = sock.accept()[0]
+        connection.settimeout(5)
+
+        while True:
+            try:
+                connection.send(connection.recv(1024))
+            except socket.timeout:
+                break
+
+        connection.close()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/test/test_line_join.vader b/test/test_line_join.vader
index 25cefbcf..9356a2b7 100644
--- a/test/test_line_join.vader
+++ b/test/test_line_join.vader
@@ -18,67 +18,67 @@ After:
   delfunction RawCallback
 
 Execute (ALE should handle empty Lists for the lines):
-  let g:last_line = ale#job#JoinNeovimOutput(1, '', [], 'nl', function('LineCallback'))
+  let g:last_line = ale#util#JoinNeovimOutput(1, '', [], 'nl', function('LineCallback'))
 
   AssertEqual [], g:lines
   AssertEqual '', g:last_line
 
 Execute (ALE should pass on full lines for NeoVim):
-  let g:last_line = ale#job#JoinNeovimOutput(1, '', ['x', 'y', ''], 'nl', function('LineCallback'))
+  let g:last_line = ale#util#JoinNeovimOutput(1, '', ['x', 'y', ''], 'nl', function('LineCallback'))
 
   AssertEqual ['x', 'y'], g:lines
   AssertEqual '', g:last_line
 
 Execute (ALE should pass on a single long line):
-  let g:last_line = ale#job#JoinNeovimOutput(1, '', ['x'], 'nl', function('LineCallback'))
+  let g:last_line = ale#util#JoinNeovimOutput(1, '', ['x'], 'nl', function('LineCallback'))
 
   AssertEqual [], g:lines
   AssertEqual 'x', g:last_line
 
 Execute (ALE should handle just a single line of output):
-  let g:last_line = ale#job#JoinNeovimOutput(1, '', ['x', ''], 'nl', function('LineCallback'))
+  let g:last_line = ale#util#JoinNeovimOutput(1, '', ['x', ''], 'nl', function('LineCallback'))
 
   AssertEqual ['x'], g:lines
   AssertEqual '', g:last_line
 
 Execute (ALE should join two incomplete pieces of large lines together):
-  let g:last_line = ale#job#JoinNeovimOutput(1, 'x', ['y'], 'nl', function('LineCallback'))
+  let g:last_line = ale#util#JoinNeovimOutput(1, 'x', ['y'], 'nl', function('LineCallback'))
 
   AssertEqual [], g:lines
   AssertEqual 'xy', g:last_line
 
 Execute (ALE join incomplete lines, and set new ones):
-  let g:last_line = ale#job#JoinNeovimOutput(1, 'x', ['y', 'z', 'a'], 'nl', function('LineCallback'))
+  let g:last_line = ale#util#JoinNeovimOutput(1, 'x', ['y', 'z', 'a'], 'nl', function('LineCallback'))
 
   AssertEqual ['xy', 'z'], g:lines
   AssertEqual 'a', g:last_line
 
 Execute (ALE join incomplete lines, and set new ones, with two elements):
-  let g:last_line = ale#job#JoinNeovimOutput(1, 'x', ['y', 'z'], 'nl', function('LineCallback'))
+  let g:last_line = ale#util#JoinNeovimOutput(1, 'x', ['y', 'z'], 'nl', function('LineCallback'))
 
   AssertEqual ['xy'], g:lines
   AssertEqual 'z', g:last_line
 
 Execute (ALE should pass on full lines for NeoVim for raw data):
-  let g:last_line = ale#job#JoinNeovimOutput(1, '', ['x', 'y', ''], 'raw', function('RawCallback'))
+  let g:last_line = ale#util#JoinNeovimOutput(1, '', ['x', 'y', ''], 'raw', function('RawCallback'))
 
   AssertEqual "x\ny\n", g:data
   AssertEqual '', g:last_line
 
 Execute (ALE should pass on a single long line):
-  let g:last_line = ale#job#JoinNeovimOutput(1, '', ['x'], 'raw', function('RawCallback'))
+  let g:last_line = ale#util#JoinNeovimOutput(1, '', ['x'], 'raw', function('RawCallback'))
 
   AssertEqual 'x', g:data
   AssertEqual '', g:last_line
 
 Execute (ALE should handle just a single line of output):
-  let g:last_line = ale#job#JoinNeovimOutput(1, '', ['x', ''], 'raw', function('RawCallback'))
+  let g:last_line = ale#util#JoinNeovimOutput(1, '', ['x', ''], 'raw', function('RawCallback'))
 
   AssertEqual "x\n", g:data
   AssertEqual '', g:last_line
 
 Execute (ALE should pass on two lines and one incomplete one):
-  let g:last_line = ale#job#JoinNeovimOutput(1, '', ['y', 'z', 'a'], 'raw', function('RawCallback'))
+  let g:last_line = ale#util#JoinNeovimOutput(1, '', ['y', 'z', 'a'], 'raw', function('RawCallback'))
 
   AssertEqual "y\nz\na", g:data
   AssertEqual '', g:last_line
diff --git a/test/test_socket_connections.vader b/test/test_socket_connections.vader
new file mode 100644
index 00000000..71a1728b
--- /dev/null
+++ b/test/test_socket_connections.vader
@@ -0,0 +1,90 @@
+Before:
+  let g:can_run_socket_tests = !has('win32')
+  \ && (exists('*ch_close') || exists('*chanclose'))
+
+  if g:can_run_socket_tests
+    call ale#test#SetDirectory('/testplugin/test')
+
+    let g:channel_id_received = 0
+    let g:data_received = ''
+
+    function! WaitForData(expected_data, timeout) abort
+      let l:ticks = 0
+
+      while l:ticks < a:timeout
+        " Sleep first, so we can switch to the callback.
+        let l:ticks += 10
+        sleep 10ms
+
+        if g:data_received is# a:expected_data
+          break
+        endif
+      endwhile
+    endfunction
+
+    function! TestCallback(channel_id, data) abort
+      let g:channel_id_received = a:channel_id
+      let g:data_received .= a:data
+    endfunction
+
+    let g:port = 10347
+    let g:pid = str2nr(system(
+    \ 'python'
+    \ . ' ' . ale#Escape(g:dir . '/dumb_tcp_server.py')
+    \ . ' ' . g:port
+    \))
+  endif
+
+After:
+  if g:can_run_socket_tests
+    call ale#test#RestoreDirectory()
+
+    unlet! g:channel_id_received
+    unlet! g:data_received
+    unlet! g:channel_id
+
+    delfunction WaitForData
+    delfunction TestCallback
+
+    if has_key(g:, 'pid')
+      call system('kill ' . g:pid)
+    endif
+
+    unlet! g:pid
+    unlet! g:port
+  endif
+
+  unlet! g:can_run_socket_tests
+
+Execute(Sending and receiving connections to sockets should work):
+  if g:can_run_socket_tests
+    let g:channel_id = ale#socket#Open(
+    \ '127.0.0.1:' . g:port,
+    \ {'callback': function('TestCallback')}
+    \)
+
+    Assert g:channel_id >= 0, 'The socket was not opened!'
+
+    call ale#socket#Send(g:channel_id, 'hello')
+    call ale#socket#Send(g:channel_id, ' world')
+
+    AssertEqual 1, ale#socket#IsOpen(g:channel_id)
+
+    " Wait up to 1 second for the expected data to arrive.
+    call WaitForData('hello world', 1000)
+
+    AssertEqual g:channel_id, g:channel_id_received
+    AssertEqual 'hello world', g:data_received
+
+    call ale#socket#Close(g:channel_id)
+
+    AssertEqual 0, ale#socket#IsOpen(g:channel_id)
+  endif
+
+  " NeoVim versions which can't connect to sockets should just fail.
+  if has('nvim') && !exists('*chanclose')
+    AssertEqual -1, ale#socket#Open(
+    \ '127.0.0.1:1111',
+    \ {'callback': function('function')}
+    \)
+  endif