Merge branch 'jm-remove-over-provisioned-watermark' into 'master'

Update zoekt watermark level whenever reservations change

See merge request https://gitlab.com/gitlab-org/gitlab/-/merge_requests/172790



Merged-by: Madelein van Niekerk <mvanniekerk@gitlab.com>
Approved-by: Madelein van Niekerk <mvanniekerk@gitlab.com>
Reviewed-by: Madelein van Niekerk <mvanniekerk@gitlab.com>
Co-authored-by: John Mason <jmason@gitlab.com>
Co-authored-by: John Mason <9717668-johnmason@users.noreply.gitlab.com>

ee/app/models/search/zoekt/index.rb

@@ -132,7 +132,9 @@ def update_reserved_storage_bytes!
 
         return if new_reserved_bytes == reserved_storage_bytes
 
-        update_column(:reserved_storage_bytes, new_reserved_bytes)
+        self.reserved_storage_bytes = new_reserved_bytes
+        self.watermark_level = appropriate_watermark_level
+        save!
       rescue StandardError => err
         logger.error(build_structured_payload(
           message: 'Error attempting to update reserved_storage_bytes',
@@ -149,12 +151,26 @@ def free_storage_bytes
         reserved_storage_bytes.to_i - used_storage_bytes
       end
 
+      def storage_percent_used
+        used_storage_bytes / reserved_storage_bytes.to_f
+      end
+
       def should_be_deleted?
         SHOULD_BE_DELETED_STATES.include? state.to_sym
       end
 
       private
 
+      def appropriate_watermark_level
+        case storage_percent_used
+        when 0...STORAGE_IDEAL_PERCENT_USED then :overprovisioned
+        when STORAGE_IDEAL_PERCENT_USED...STORAGE_LOW_WATERMARK then :healthy
+        when STORAGE_LOW_WATERMARK...STORAGE_HIGH_WATERMARK then :low_watermark_exceeded
+        else
+          :high_watermark_exceeded
+        end
+      end
+
       def delete_from_index
         ::Search::Zoekt::NamespaceIndexerWorker.perform_async(namespace_id, 'delete', zoekt_node_id)
       end

ee/app/workers/search/zoekt/index_watermark_changed_event_worker.rb

@@ -17,9 +17,7 @@ def handle_event(event)
         return unless watermark_level.present? && index_ids.present?
 
         Search::Zoekt::Index.id_in(event.data[:index_ids]).each_batch do |batch|
-          batch.update_all(watermark_level: event.data[:watermark_level])
           batch.each(&:update_reserved_storage_bytes!)
-          log_extra_metadata_on_done(watermark_level, index_ids.length)
         end
       end
     end

ee/spec/models/search/zoekt/index_spec.rb

@@ -282,6 +282,57 @@
       ).to eq(described_class::STORAGE_IDEAL_PERCENT_USED)
     end
 
+    describe 'updates the watermark level to the appropriate state' do
+      let(:percent_used) { 0 }
+
+      before do
+        idx.high_watermark_exceeded!
+        allow(idx).to receive(:storage_percent_used).and_return(percent_used)
+      end
+
+      context 'when index should be marked as overprovisioned' do
+        let(:percent_used) { 0.01 }
+
+        it 'updates the watermark level to the appropriate state' do
+          expect { idx.update_reserved_storage_bytes! }.to change {
+            idx.reload.watermark_level
+          }.from("high_watermark_exceeded").to("overprovisioned")
+        end
+      end
+
+      context 'when index should be marked as healthy' do
+        let(:percent_used) { described_class::STORAGE_IDEAL_PERCENT_USED }
+
+        it 'updates the watermark level to the appropriate state' do
+          expect { idx.update_reserved_storage_bytes! }.to change {
+            idx.reload.watermark_level
+          }.from("high_watermark_exceeded").to("healthy")
+        end
+      end
+
+      context 'when index should be marked as low_watermark_exceeded' do
+        let(:percent_used) { described_class::STORAGE_LOW_WATERMARK }
+
+        it 'updates the watermark level to the appropriate state' do
+          expect { idx.update_reserved_storage_bytes! }.to change {
+            idx.reload.watermark_level
+          }.from("high_watermark_exceeded").to("low_watermark_exceeded")
+        end
+      end
+
+      context 'when index should be marked as high_watermark_exceeded' do
+        let(:percent_used) { described_class::STORAGE_HIGH_WATERMARK }
+
+        it 'updates the watermark level to the appropriate state' do
+          idx.low_watermark_exceeded!
+
+          expect { idx.update_reserved_storage_bytes! }.to change {
+            idx.reload.watermark_level
+          }.from("low_watermark_exceeded").to("high_watermark_exceeded")
+        end
+      end
+    end
+
     context 'when the node only has a little bit more storage' do
       let_it_be(:zoekt_node) { create(:zoekt_node, total_bytes: 102, used_bytes: 0) }
 

ee/spec/workers/search/zoekt/index_watermark_changed_event_worker_spec.rb

@@ -6,19 +6,9 @@
   let(:event) { Search::Zoekt::IndexWatermarkChangedEvent.new(data: data) }
 
   let_it_be(:watermark_level) { 'low_watermark_exceeded' }
-  let_it_be(:idx) { create(:zoekt_index, reserved_storage_bytes: max_storage_bytes) }
-  let_it_be(:idx_2) { create(:zoekt_index) }
-  let_it_be(:index_ids) { [idx.id] }
+  let_it_be(:indices_in_event) { create_list(:zoekt_index, 3) }
+  let_it_be(:index_ids) { indices_in_event.map(&:id) }
   let_it_be(:max_storage_bytes) { 100 }
-  let_it_be(:idx_project) { create(:project, namespace_id: idx.namespace_id) }
-  let_it_be(:idx_project_2) { create(:project, namespace_id: idx_2.namespace_id) }
-  let_it_be(:repo) do
-    idx.zoekt_repositories.create!(zoekt_index: idx, project: idx_project, state: :ready)
-  end
-
-  let_it_be(:repo_2) do
-    idx.zoekt_repositories.create!(zoekt_index: idx_2, project: idx_project_2, state: :ready, size_bytes: 1)
-  end
 
   let(:data) do
     { index_ids: index_ids, watermark_level: watermark_level }
@@ -27,22 +17,15 @@
   it_behaves_like 'subscribes to event'
 
   it_behaves_like 'an idempotent worker' do
-    context 'when there is an index that is over low storage watermark' do
-      it 'updates the index to be watermarked' do
-        expect do
-          consume_event(subscriber: described_class, event: event)
-        end.to change { Search::Zoekt::Index.low_watermark_exceeded.count }.from(0).to(1)
-      end
-    end
-
-    context 'when there is an index that is over a storage watermark' do
-      let_it_be(:watermark_level) { 'high_watermark_exceeded' }
+    it 'calls `update_reserved_storage_bytes` for each index' do
+      expect(::Search::Zoekt::Index).to receive(:id_in).with(index_ids).and_return(indices_in_event)
+      expect(indices_in_event).to receive(:each_batch).and_yield(indices_in_event)
 
-      it 'updates the index to be watermarked' do
-        expect do
-          consume_event(subscriber: described_class, event: event)
-        end.to change { Search::Zoekt::Index.high_watermark_exceeded.count }.from(0).to(1)
+      indices_in_event.each do |zkt_index| # rubocop:disable RSpec/IteratedExpectation -- Not applicable
+        expect(zkt_index).to receive(:update_reserved_storage_bytes!)
       end
+
+      consume_event(subscriber: described_class, event: event)
     end
   end
 end