// 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.progress.ProgressManager;
import com.intellij.openapi.progress.util.ProgressIndicatorUtils;
import org.jetbrains.annotations.NotNull;

import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

final class StorageGuard {
  private final Lock myLock = new ReentrantLock();
  private final Condition myCondition = myLock.newCondition();

  private int myHolds;

  public interface StorageModeExitHandler {
    void leave();
  }

  private final StorageModeExitHandler myTrueStorageModeExitHandler = new StorageModeExitHandler() {
    @Override
    public void leave() {
      StorageGuard.this.leave(true);
    }
  };
  private final StorageModeExitHandler myFalseStorageModeExitHandler = new StorageModeExitHandler() {
    @Override
    public void leave() {
      StorageGuard.this.leave(false);
    }
  };

  @SuppressWarnings("WhileLoopSpinsOnField")
  @NotNull
  StorageModeExitHandler enter(boolean mode) {
    ProgressManager progressManager = ProgressManager.getInstance();
    assert !progressManager.isInNonCancelableSection();
    myLock.lock();
    try {
      if (mode) {
        while (myHolds < 0) {
          ProgressIndicatorUtils.awaitWithCheckCanceled(myCondition);
        }
        myHolds++;
        return myTrueStorageModeExitHandler;
      }
      else {
        while (myHolds > 0) {
          ProgressIndicatorUtils.awaitWithCheckCanceled(myCondition);
        }
        myHolds--;
        return myFalseStorageModeExitHandler;
      }
    }
    finally {
      myLock.unlock();
    }
  }

  private void leave(boolean mode) {
    myLock.lock();
    try {
      myHolds += mode ? -1 : 1;
      if (myHolds == 0) {
        myCondition.signalAll();
      }
    }
    finally {
      myLock.unlock();
    }
  }
}

@@ -1194,39 +1194,36 @@
     
-ProgressManager progressManager
+boolean nonCancelableSection
  = P
@@ -1254,37 +1254,8 @@
 ce()
-;%0A    assert !progressManager
 .isI
@@ -1370,61 +1370,30 @@
     
-ProgressIndicatorUtils.awaitWithCheckCanceled(myCondi
+await(nonCancelableSec
 tion
@@ -1534,61 +1534,30 @@
     
-ProgressIndicatorUtils.awaitWithCheckCanceled(myCondi
+await(nonCancelableSec
 tion
@@ -1654,32 +1654,32 @@
 %7D%0A    finally %7B%0A
-
       myLock.unl
@@ -1692,24 +1692,245 @@
 %0A    %7D%0A  %7D%0A%0A
+  private void await(boolean nonCancelableSection) %7B%0A    if (nonCancelableSection) %7B%0A      myCondition.awaitUninterruptibly();%0A    %7D%0A    else %7B%0A      ProgressIndicatorUtils.awaitWithCheckCanceled(myCondition);%0A    %7D%0A  %7D%0A%0A
   private vo

add different await mods for StorageGuard