BigW Consortium Gitlab

gitaly_client.rb 2.97 KB
Newer Older
1 2
require 'base64'

3 4 5 6
require 'gitaly'

module Gitlab
  module GitalyClient
7 8 9 10 11 12
    module MigrationStatus
      DISABLED = 1
      OPT_IN = 2
      OPT_OUT = 3
    end

13 14
    SERVER_VERSION_FILE = 'GITALY_SERVER_VERSION'.freeze

15 16
    MUTEX = Mutex.new
    private_constant :MUTEX
17

18 19 20 21 22 23 24 25 26
    def self.stub(name, storage)
      MUTEX.synchronize do
        @stubs ||= {}
        @stubs[storage] ||= {}
        @stubs[storage][name] ||= begin
          klass = Gitaly.const_get(name.to_s.camelcase.to_sym).const_get(:Stub)
          addr = address(storage)
          addr = addr.sub(%r{^tcp://}, '') if URI(addr).scheme == 'tcp'
          klass.new(addr, :this_channel_is_insecure)
27 28
        end
      end
29 30
    end

31 32 33 34
    def self.clear_stubs!
      MUTEX.synchronize do
        @stubs = nil
      end
35 36
    end

37 38 39
    def self.address(storage)
      params = Gitlab.config.repositories.storages[storage]
      raise "storage not found: #{storage.inspect}" if params.nil?
40

41 42 43 44
      address = params['gitaly_address']
      unless address.present?
        raise "storage #{storage.inspect} is missing a gitaly_address"
      end
45

46 47
      unless URI(address).scheme.in?(%w(tcp unix))
        raise "Unsupported Gitaly address: #{address.inspect} does not use URL scheme 'tcp' or 'unix'"
48 49
      end

50
      address
51 52
    end

53 54
    # All Gitaly RPC call sites should use GitalyClient.call. This method
    # makes sure that per-request authentication headers are set.
55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72
    def self.call(storage, service, rpc, request)
      metadata = request_metadata(storage)
      metadata = yield(metadata) if block_given?
      stub(service, storage).send(rpc, request, metadata)
    end
  
    def self.request_metadata(storage)
      encoded_token = Base64.strict_encode64(token(storage).to_s)
      { metadata: { 'authorization' => "Bearer #{encoded_token}" } }
    end

    def self.token(storage)
      params = Gitlab.config.repositories.storages[storage]
      raise "storage not found: #{storage.inspect}" if params.nil?

      params['gitaly_token'].presence || Gitlab.config.gitaly['token']
    end

73
    def self.enabled?
74
      Gitlab.config.gitaly.enabled
75
    end
76

77 78 79 80 81 82 83 84 85 86 87 88 89 90
    def self.feature_enabled?(feature, status: MigrationStatus::OPT_IN)
      return false if !enabled? || status == MigrationStatus::DISABLED

      feature = Feature.get("gitaly_#{feature}")

      # If the feature hasn't been set, turn it on if it's opt-out
      return status == MigrationStatus::OPT_OUT unless Feature.persisted?(feature)

      if feature.percentage_of_time_value > 0
        # Probabilistically enable this feature
        return Random.rand() * 100 < feature.percentage_of_time_value
      end

      feature.enabled?
91 92 93 94 95 96 97 98 99 100 101
    end

    def self.migrate(feature)
      is_enabled  = feature_enabled?(feature)
      metric_name = feature.to_s
      metric_name += "_gitaly" if is_enabled

      Gitlab::Metrics.measure(metric_name) do
        yield is_enabled
      end
    end
102 103 104 105 106

    def self.expected_server_version
      path = Rails.root.join(SERVER_VERSION_FILE)
      path.read.chomp
    end
107 108
  end
end