Add a gitlab:tcp_check rake task

changelogs/unreleased/add-tcp-check-rake-task.yml

+---
+title: Add a gitlab:tcp_check rake task
+merge_request: 15759
+author:
+type: added

doc/administration/raketasks/maintenance.md

@@ -221,3 +221,22 @@ sudo gitlab-rake gitlab:shell:create_hooks
 cd /home/git/gitlab
 sudo -u git -H bundle exec rake gitlab:shell:create_hooks RAILS_ENV=production
 ```
+
+## Check TCP connectivity to a remote site
+
+Sometimes you need to know if your GitLab installation can connect to a TCP
+service on another machine - perhaps a PostgreSQL or HTTPS server. A rake task
+is included to help you with this:
+
+**Omnibus Installation**
+
+```
+sudo gitlab-rake gitlab:tcp_check[example.com,80]
+```
+
+**Source Installation**
+
+```
+cd /home/git/gitlab
+sudo -u git -H bundle exec rake gitlab:tcp_check[example.com,80] RAILS_ENV=production
+```

lib/gitlab/tcp_checker.rb

+module Gitlab
+  class TcpChecker
+    attr_reader :remote_host, :remote_port, :local_host, :local_port, :error
+
+    def initialize(remote_host, remote_port, local_host = nil, local_port = nil)
+      @remote_host = remote_host
+      @remote_port = remote_port
+      @local_host = local_host
+      @local_port = local_port
+    end
+
+    def local
+      join_host_port(local_host, local_port)
+    end
+
+    def remote
+      join_host_port(remote_host, remote_port)
+    end
+
+    def check(timeout: 10)
+      Socket.tcp(
+        remote_host, remote_port,
+        local_host, local_port,
+        connect_timeout: timeout
+      ) do |sock|
+        @local_port, @local_host = Socket.unpack_sockaddr_in(sock.local_address)
+        @remote_port, @remote_host = Socket.unpack_sockaddr_in(sock.remote_address)
+      end
+
+      true
+    rescue => err
+      @error = err
+
+      false
+    end
+
+    private
+
+    def join_host_port(host, port)
+      host = "[#{host}]" if host.include?(':')
+
+      "#{host}:#{port}"
+    end
+  end
+end

lib/tasks/gitlab/tcp_check.rake

+namespace :gitlab do
+  desc "GitLab | Check TCP connectivity to a specific host and port"
+  task :tcp_check, [:host, :port] => :environment do |_t, args|
+    unless args.host && args.port
+      puts "Please specify a host and port: `rake gitlab:tcp_check[example.com,80]`".color(:red)
+      exit 1
+    end
+
+    checker = Gitlab::TcpChecker.new(args.host, args.port)
+
+    if checker.check
+      puts "TCP connection from #{checker.local} to #{checker.remote} succeeded".color(:green)
+    else
+      puts "TCP connection to #{checker.remote} failed: #{checker.error}".color(:red)
+      puts
+      puts 'Check that host and port are correct, and that the traffic is permitted through any firewalls.'
+      exit 1
+    end
+  end
+end

spec/lib/gitlab/tcp_checker_spec.rb

+require 'spec_helper'
+
+describe Gitlab::TcpChecker do
+  before do
+    @server = TCPServer.new('localhost', 0)
+    _, @port, _, @ip = @server.addr
+  end
+
+  after do
+    @server.close
+  end
+
+  subject(:checker) { described_class.new(@ip, @port) }
+
+  describe '#check' do
+    subject { checker.check }
+
+    it 'can connect to an open port' do
+      is_expected.to be_truthy
+
+      expect(checker.error).to be_nil
+    end
+
+    it 'fails to connect to a closed port' do
+      @server.close
+
+      is_expected.to be_falsy
+
+      expect(checker.error).to be_a(Errno::ECONNREFUSED)
+    end
+  end
+end