// Copyright 2000-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
package com.intellij.util.indexing;

import com.intellij.openapi.diagnostic.Logger;
import com.intellij.openapi.util.Computable;
import org.jetbrains.annotations.NotNull;

import java.util.stream.Stream;

abstract class StorageBufferingHandler {
  private static final Logger LOG = Logger.getInstance(StorageBufferingHandler.class);

  private final StorageGuard myStorageLock = new StorageGuard();
  private volatile boolean myPreviousDataBufferingState;
  private final Object myBufferingStateUpdateLock = new Object();

  boolean runUpdate(boolean transientInMemoryIndices, @NotNull Computable<Boolean> update) {
    StorageGuard.StorageModeExitHandler storageModeExitHandler = myStorageLock.enter(transientInMemoryIndices);

    if (myPreviousDataBufferingState != transientInMemoryIndices) {
      synchronized (myBufferingStateUpdateLock) {
        if (myPreviousDataBufferingState != transientInMemoryIndices) {
          getIndexes().forEach(index -> {
            assert index != null;
            index.setBufferingEnabled(transientInMemoryIndices);
          });
          myPreviousDataBufferingState = transientInMemoryIndices;
        }
      }
    }

    try {
      return update.compute();
    } finally {
      storageModeExitHandler.leave();
    }
  }

  void resetState() {
    myPreviousDataBufferingState = false;
  }

  void assertTransientMode() {
    LOG.assertTrue(myPreviousDataBufferingState);
  }

  void assertOnTheDiskMode() {
    LOG.assertTrue(!myPreviousDataBufferingState);
  }

  @NotNull
  protected abstract Stream<UpdatableIndex<?, ? ,?>> getIndexes();
}

@@ -862,16 +862,247 @@
 dices);%0A
+    try %7B%0A      ensureBufferingState(transientInMemoryIndices);%0A      return update.compute();%0A    %7D finally %7B%0A      storageModeExitHandler.leave();%0A    %7D%0A  %7D%0A%0A  private void ensureBufferingState(boolean transientInMemoryIndices) %7B
 %0A    if 
@@ -1533,110 +1533,8 @@
    %7D
-%0A%0A    try %7B%0A      return update.compute();%0A    %7D finally %7B%0A      storageModeExitHandler.leave();%0A    %7D
 %0A  %7D

change index buffering state inside try-finally to ensure we exit from storage guard in case of PCE or other exceptions