1# frozen_string_literal: true
2
3module Gitlab
4  module Database
5    module LoadBalancing
6      # Module used for handling sticking connections to a primary, if
7      # necessary.
8      class Sticking
9        # The number of seconds after which a session should stop reading from
10        # the primary.
11        EXPIRATION = 30
12
13        def initialize(load_balancer)
14          @load_balancer = load_balancer
15        end
16
17        # Unsticks or continues sticking the current request.
18        #
19        # This method also updates the Rack environment so #call can later
20        # determine if we still need to stick or not.
21        #
22        # env - The Rack environment.
23        # namespace - The namespace to use for sticking.
24        # id - The identifier to use for sticking.
25        # model - The ActiveRecord model to scope sticking to.
26        def stick_or_unstick_request(env, namespace, id)
27          unstick_or_continue_sticking(namespace, id)
28
29          env[::Gitlab::Database::LoadBalancing::RackMiddleware::STICK_OBJECT] ||= Set.new
30          env[::Gitlab::Database::LoadBalancing::RackMiddleware::STICK_OBJECT] << [self, namespace, id]
31        end
32
33        # Sticks to the primary if a write was performed.
34        def stick_if_necessary(namespace, id)
35          stick(namespace, id) if ::Gitlab::Database::LoadBalancing::Session.current.performed_write?
36        end
37
38        def all_caught_up?(namespace, id)
39          location = last_write_location_for(namespace, id)
40
41          return true unless location
42
43          @load_balancer.select_up_to_date_host(location).tap do |found|
44            ActiveSupport::Notifications.instrument(
45              'caught_up_replica_pick.load_balancing',
46              { result: found }
47            )
48
49            unstick(namespace, id) if found
50          end
51        end
52
53        # Selects hosts that have caught up with the primary. This ensures
54        # atomic selection of the host to prevent the host list changing
55        # in another thread.
56        #
57        # Returns true if one host was selected.
58        def select_caught_up_replicas(namespace, id)
59          location = last_write_location_for(namespace, id)
60
61          # Unlike all_caught_up?, we return false if no write location exists.
62          # We want to be sure we talk to a replica that has caught up for a specific
63          # write location. If no such location exists, err on the side of caution.
64          return false unless location
65
66          @load_balancer.select_up_to_date_host(location).tap do |selected|
67            unstick(namespace, id) if selected
68          end
69        end
70
71        # Sticks to the primary if necessary, otherwise unsticks an object (if
72        # it was previously stuck to the primary).
73        def unstick_or_continue_sticking(namespace, id)
74          return if all_caught_up?(namespace, id)
75
76          ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
77        end
78
79        # Select a replica that has caught up with the primary. If one has not been
80        # found, stick to the primary.
81        def select_valid_host(namespace, id)
82          replica_selected =
83            select_caught_up_replicas(namespace, id)
84
85          ::Gitlab::Database::LoadBalancing::Session.current.use_primary! unless replica_selected
86        end
87
88        # Starts sticking to the primary for the given namespace and id, using
89        # the latest WAL pointer from the primary.
90        def stick(namespace, id)
91          mark_primary_write_location(namespace, id)
92          ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
93        end
94
95        def bulk_stick(namespace, ids)
96          with_primary_write_location do |location|
97            ids.each do |id|
98              set_write_location_for(namespace, id, location)
99            end
100          end
101
102          ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
103        end
104
105        def with_primary_write_location
106          # When only using the primary, there's no point in getting write
107          # locations, as the primary is always in sync with itself.
108          return if @load_balancer.primary_only?
109
110          location = @load_balancer.primary_write_location
111
112          return if location.blank?
113
114          yield(location)
115        end
116
117        def mark_primary_write_location(namespace, id)
118          with_primary_write_location do |location|
119            set_write_location_for(namespace, id, location)
120          end
121        end
122
123        def unstick(namespace, id)
124          Gitlab::Redis::SharedState.with do |redis|
125            redis.del(redis_key_for(namespace, id))
126          end
127        end
128
129        def set_write_location_for(namespace, id, location)
130          Gitlab::Redis::SharedState.with do |redis|
131            redis.set(redis_key_for(namespace, id), location, ex: EXPIRATION)
132          end
133        end
134
135        def last_write_location_for(namespace, id)
136          Gitlab::Redis::SharedState.with do |redis|
137            redis.get(redis_key_for(namespace, id))
138          end
139        end
140
141        def redis_key_for(namespace, id)
142          name = @load_balancer.name
143
144          "database-load-balancing/write-location/#{name}/#{namespace}/#{id}"
145        end
146      end
147    end
148  end
149end
150