src/main/java/com/gerritforge/gerrit/globalrefdb/validation/BatchRefUpdateValidator.java - modules/global-refdb - Git at Google

 // Copyright (C) 2020 The Android Open Source Project
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 // http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 package com.gerritforge.gerrit.globalrefdb.validation;

 import com.gerritforge.gerrit.globalrefdb.validation.dfsrefdb.CustomSharedRefEnforcementByProject;
 import com.gerritforge.gerrit.globalrefdb.validation.dfsrefdb.DefaultSharedRefEnforcement;
 import com.gerritforge.gerrit.globalrefdb.validation.dfsrefdb.OutOfSyncException;
 import com.gerritforge.gerrit.globalrefdb.validation.dfsrefdb.SharedRefEnforcement;
 import com.gerritforge.gerrit.globalrefdb.validation.dfsrefdb.SharedRefEnforcement.EnforcePolicy;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.flogger.FluentLogger;
 import com.google.inject.Inject;
 import com.google.inject.assistedinject.Assisted;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import org.eclipse.jgit.lib.BatchRefUpdate;
 import org.eclipse.jgit.lib.ObjectId;
 import org.eclipse.jgit.lib.ObjectIdRef;
 import org.eclipse.jgit.lib.Ref;
 import org.eclipse.jgit.lib.RefDatabase;
 import org.eclipse.jgit.transport.ReceiveCommand;

 /** Enables the detection of out-of-sync by validating batch ref updates against the global refdb */
 public class BatchRefUpdateValidator extends RefUpdateValidator {
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();

   /** {@code BatchRefUpdateValidator} Factory for Guice assisted injection. */
   public interface Factory {
     BatchRefUpdateValidator create(
         String projectName, RefDatabase refDb, ImmutableSet<String> ignoredRefs);
   }

   public interface BatchValidationWrapper {
     void apply(BatchRefUpdate batchRefUpdate, NoParameterVoidFunction arg) throws IOException;
   }

   /**
    * Constructs a {@code BatchRefUpdateValidator} able to check the validity of batch ref-updates
    * against global refdb before execution.
    *
    * @param sharedRefDb an instance of the global refdb to check for out-of-sync refs.
    * @param validationMetrics to update validation results, such as split-brains.
    * @param refEnforcement Specific ref enforcements for this project. Either a {@link
    *     CustomSharedRefEnforcementByProject} when custom policies are provided via configuration *
    *     file or a {@link DefaultSharedRefEnforcement} for defaults.
    * @param lockWrapperFactory factory providing a {@link LockWrapper}
    * @param projectsFilter filter to match whether the project being updated should be validated
    *     against global refdb
    * @param projectName the name of the project being updated.
    * @param refDb for ref operations
    * @param ignoredRefs A set of refs for which updates should not be checked against the shared
    *     ref-db
    */
   @Inject
   public BatchRefUpdateValidator(
       SharedRefDatabaseWrapper sharedRefDb,
       ValidationMetrics validationMetrics,
       SharedRefEnforcement refEnforcement,
       LockWrapper.Factory lockWrapperFactory,
       ProjectsFilter projectsFilter,
       @Assisted String projectName,
       @Assisted RefDatabase refDb,
       @Assisted ImmutableSet<String> ignoredRefs) {
     super(
         sharedRefDb,
         validationMetrics,
         refEnforcement,
         lockWrapperFactory,
         projectsFilter,
         projectName,
         refDb,
         ignoredRefs);
   }

   /**
    * Checks whether the provided batchRefUpdate should be validated first against the global refdb.
    * If not it just execute the provided batchRefUpdateFunction. Upon success the batchRefUpdate is
    * returned, upon failure split brain metrics are incremented and a {@link IOException} is thrown.
    *
    * <p>Validation is performed when either of these condition is true:
    *
    * <ul>
    *   <li>The project being updated is a global project ({@link
    *       RefUpdateValidator#isGlobalProject(String)}
    *   <li>The enforcement policy for the project being updated is {@link EnforcePolicy#IGNORED}
    * </ul>
    *
    * @param batchRefUpdate batchRefUpdate object
    * @param batchRefUpdateFunction batchRefUpdate function to execute upon validation
    * @param batchRefUpdateRollbackFunction function to invoke when the ref-update needs to be rolled
    *     back
    * @throws IOException batch update failed
    */
   @SuppressWarnings("JavadocReference")
   public void executeBatchUpdateWithValidation(
       BatchRefUpdate batchRefUpdate,
       NoParameterVoidFunction batchRefUpdateFunction,
       OneParameterVoidFunction<List<ReceiveCommand>> batchRefUpdateRollbackFunction)
       throws IOException {
     if (refEnforcement.getPolicy(projectName) == EnforcePolicy.IGNORED
         || !isGlobalProject(projectName)) {
       batchRefUpdateFunction.invoke();
       return;
     }

     try {
       doExecuteBatchUpdate(batchRefUpdate, batchRefUpdateFunction, batchRefUpdateRollbackFunction);
     } catch (IOException e) {
       logger.atWarning().withCause(e).log(
           "Failed to execute Batch Update on project %s", projectName);
       if (refEnforcement.getPolicy(projectName) == EnforcePolicy.REQUIRED) {
         throw e;
       }
     }
   }

   private void doExecuteBatchUpdate(
       BatchRefUpdate batchRefUpdate,
       NoParameterVoidFunction delegateUpdate,
       OneParameterVoidFunction<List<ReceiveCommand>> delegateUpdateRollback)
       throws IOException {

     List<ReceiveCommand> commands = batchRefUpdate.getCommands();
     if (commands.isEmpty()) {
       return;
     }

     List<RefPair> refsToUpdate = getRefPairs(commands).collect(Collectors.toList());
     List<RefPair> refsFailures =
         refsToUpdate.stream().filter(RefPair::hasFailed).collect(Collectors.toList());
     if (!refsFailures.isEmpty()) {
       String allFailuresMessage =
           refsFailures.stream()
               .map(refPair -> String.format("Failed to fetch ref %s", refPair.compareRef.getName()))
               .collect(Collectors.joining(", "));
       Exception firstFailureException = refsFailures.get(0).exception;

       logger.atSevere().withCause(firstFailureException).log(allFailuresMessage);
       throw new IOException(allFailuresMessage, firstFailureException);
     }

     try (CloseableSet<AutoCloseable> locks = new CloseableSet<>()) {
       final List<RefPair> finalRefsToUpdate = compareAndGetLatestLocalRefs(refsToUpdate, locks);
       delegateUpdate.invoke();
       boolean sharedDbUpdateSucceeded = false;
       try {
         updateSharedRefDb(batchRefUpdate.getCommands().stream(), finalRefsToUpdate);
         sharedDbUpdateSucceeded = true;
       } catch (Exception e) {
         logger.atWarning().withCause(e).log(
             "Batch ref-update failed because of failure during the global refdb update.");
       } finally {
         if (!sharedDbUpdateSucceeded) {
           List<ReceiveCommand> receiveCommands = batchRefUpdate.getCommands();
           logger.atWarning().log(
               "Batch ref-update failed, set all commands Result to LOCK_FAILURE [%d]",
               commands.size());
           rollback(delegateUpdateRollback, finalRefsToUpdate, receiveCommands);
         }
       }
     } catch (OutOfSyncException e) {
       List<ReceiveCommand> receiveCommands = batchRefUpdate.getCommands();
       logger.atWarning().withCause(e).log(
           "Batch ref-update failing because node is out of sync with the shared ref-db. Set all commands Result to LOCK_FAILURE [%d]",
           receiveCommands.size());
       receiveCommands.forEach((command) -> command.setResult(ReceiveCommand.Result.LOCK_FAILURE));
     }
   }

   private void rollback(
       OneParameterVoidFunction<List<ReceiveCommand>> delegateUpdateRollback,
       List<RefPair> refsBeforeUpdate,
       List<ReceiveCommand> receiveCommands)
       throws IOException {
     List<ReceiveCommand> rollbackCommands =
         refsBeforeUpdate.stream()
             .map(
                 refBeforeUpdate ->
                     new ReceiveCommand(
                         refBeforeUpdate.putValue,
                         refBeforeUpdate.compareRef.getObjectId(),
                         refBeforeUpdate.getName()))
             .collect(Collectors.toList());
     delegateUpdateRollback.invoke(rollbackCommands);
     receiveCommands.forEach(command -> command.setResult(ReceiveCommand.Result.LOCK_FAILURE));
   }

   private void updateSharedRefDb(Stream<ReceiveCommand> commandStream, List<RefPair> refsToUpdate)
       throws IOException {
     if (commandStream.anyMatch(cmd -> cmd.getResult() != ReceiveCommand.Result.OK)) {
       return;
     }

     for (RefPair refPair : refsToUpdate) {
       updateSharedDbOrThrowExceptionFor(refPair);
     }
   }

   private Stream<RefPair> getRefPairs(List<ReceiveCommand> receivedCommands) {
     return receivedCommands.stream().map(this::getRefPairForCommand);
   }

   private RefPair getRefPairForCommand(ReceiveCommand command) {
     try {
       switch (command.getType()) {
         case CREATE:
           return new RefPair(nullRef(command.getRefName()), getNewRef(command));

         case UPDATE:
         case UPDATE_NONFASTFORWARD:
           return new RefPair(getCurrentRef(command.getRefName()), getNewRef(command));

         case DELETE:
           return new RefPair(getCurrentRef(command.getRefName()), ObjectId.zeroId());

         default:
           return new RefPair(
               command.getRef(),
               new IllegalArgumentException("Unsupported command type " + command.getType()));
       }
     } catch (IOException e) {
       return new RefPair(command.getRef(), e);
     }
   }

   private ObjectId getNewRef(ReceiveCommand command) {
     return command.getNewId();
   }

   private List<RefPair> compareAndGetLatestLocalRefs(
       List<RefPair> refsToUpdate, CloseableSet<AutoCloseable> locks) throws IOException {
     List<RefPair> latestRefsToUpdate = new ArrayList<>();
     for (RefPair refPair : refsToUpdate) {
       latestRefsToUpdate.add(compareAndGetLatestLocalRef(refPair, locks));
     }
     return latestRefsToUpdate;
   }

   private static final Ref nullRef(String refName) {
     return new ObjectIdRef.Unpeeled(Ref.Storage.NETWORK, refName, ObjectId.zeroId());
   }
 }
	// Copyright (C) 2020 The Android Open Source Project
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	package com.gerritforge.gerrit.globalrefdb.validation;

	import com.gerritforge.gerrit.globalrefdb.validation.dfsrefdb.CustomSharedRefEnforcementByProject;
	import com.gerritforge.gerrit.globalrefdb.validation.dfsrefdb.DefaultSharedRefEnforcement;
	import com.gerritforge.gerrit.globalrefdb.validation.dfsrefdb.OutOfSyncException;
	import com.gerritforge.gerrit.globalrefdb.validation.dfsrefdb.SharedRefEnforcement;
	import com.gerritforge.gerrit.globalrefdb.validation.dfsrefdb.SharedRefEnforcement.EnforcePolicy;
	import com.google.common.collect.ImmutableSet;
	import com.google.common.flogger.FluentLogger;
	import com.google.inject.Inject;
	import com.google.inject.assistedinject.Assisted;
	import java.io.IOException;
	import java.util.ArrayList;
	import java.util.List;
	import java.util.stream.Collectors;
	import java.util.stream.Stream;
	import org.eclipse.jgit.lib.BatchRefUpdate;
	import org.eclipse.jgit.lib.ObjectId;
	import org.eclipse.jgit.lib.ObjectIdRef;
	import org.eclipse.jgit.lib.Ref;
	import org.eclipse.jgit.lib.RefDatabase;
	import org.eclipse.jgit.transport.ReceiveCommand;

	/** Enables the detection of out-of-sync by validating batch ref updates against the global refdb */
	public class BatchRefUpdateValidator extends RefUpdateValidator {
	private static final FluentLogger logger = FluentLogger.forEnclosingClass();

	/** {@code BatchRefUpdateValidator} Factory for Guice assisted injection. */
	public interface Factory {
	BatchRefUpdateValidator create(
	String projectName, RefDatabase refDb, ImmutableSet<String> ignoredRefs);
	}

	public interface BatchValidationWrapper {
	void apply(BatchRefUpdate batchRefUpdate, NoParameterVoidFunction arg) throws IOException;
	}

	/**
	* Constructs a {@code BatchRefUpdateValidator} able to check the validity of batch ref-updates
	* against global refdb before execution.
	*
	* @param sharedRefDb an instance of the global refdb to check for out-of-sync refs.
	* @param validationMetrics to update validation results, such as split-brains.
	* @param refEnforcement Specific ref enforcements for this project. Either a {@link
	* CustomSharedRefEnforcementByProject} when custom policies are provided via configuration *
	* file or a {@link DefaultSharedRefEnforcement} for defaults.
	* @param lockWrapperFactory factory providing a {@link LockWrapper}
	* @param projectsFilter filter to match whether the project being updated should be validated
	* against global refdb
	* @param projectName the name of the project being updated.
	* @param refDb for ref operations
	* @param ignoredRefs A set of refs for which updates should not be checked against the shared
	* ref-db
	*/
	@Inject
	public BatchRefUpdateValidator(
	SharedRefDatabaseWrapper sharedRefDb,
	ValidationMetrics validationMetrics,
	SharedRefEnforcement refEnforcement,
	LockWrapper.Factory lockWrapperFactory,
	ProjectsFilter projectsFilter,
	@Assisted String projectName,
	@Assisted RefDatabase refDb,
	@Assisted ImmutableSet<String> ignoredRefs) {
	super(
	sharedRefDb,
	validationMetrics,
	refEnforcement,
	lockWrapperFactory,
	projectsFilter,
	projectName,
	refDb,
	ignoredRefs);
	}

	/**
	* Checks whether the provided batchRefUpdate should be validated first against the global refdb.
	* If not it just execute the provided batchRefUpdateFunction. Upon success the batchRefUpdate is
	* returned, upon failure split brain metrics are incremented and a {@link IOException} is thrown.
	*
	* <p>Validation is performed when either of these condition is true:
	*
	* <ul>
	* <li>The project being updated is a global project ({@link
	* RefUpdateValidator#isGlobalProject(String)}
	* <li>The enforcement policy for the project being updated is {@link EnforcePolicy#IGNORED}
	* </ul>
	*
	* @param batchRefUpdate batchRefUpdate object
	* @param batchRefUpdateFunction batchRefUpdate function to execute upon validation
	* @param batchRefUpdateRollbackFunction function to invoke when the ref-update needs to be rolled
	* back
	* @throws IOException batch update failed
	*/
	@SuppressWarnings("JavadocReference")
	public void executeBatchUpdateWithValidation(
	BatchRefUpdate batchRefUpdate,
	NoParameterVoidFunction batchRefUpdateFunction,
	OneParameterVoidFunction<List<ReceiveCommand>> batchRefUpdateRollbackFunction)
	throws IOException {
	if (refEnforcement.getPolicy(projectName) == EnforcePolicy.IGNORED
	\|\| !isGlobalProject(projectName)) {
	batchRefUpdateFunction.invoke();
	return;
	}

	try {
	doExecuteBatchUpdate(batchRefUpdate, batchRefUpdateFunction, batchRefUpdateRollbackFunction);
	} catch (IOException e) {
	logger.atWarning().withCause(e).log(
	"Failed to execute Batch Update on project %s", projectName);
	if (refEnforcement.getPolicy(projectName) == EnforcePolicy.REQUIRED) {
	throw e;
	}
	}
	}

	private void doExecuteBatchUpdate(
	BatchRefUpdate batchRefUpdate,
	NoParameterVoidFunction delegateUpdate,
	OneParameterVoidFunction<List<ReceiveCommand>> delegateUpdateRollback)
	throws IOException {

	List<ReceiveCommand> commands = batchRefUpdate.getCommands();
	if (commands.isEmpty()) {
	return;
	}

	List<RefPair> refsToUpdate = getRefPairs(commands).collect(Collectors.toList());
	List<RefPair> refsFailures =
	refsToUpdate.stream().filter(RefPair::hasFailed).collect(Collectors.toList());
	if (!refsFailures.isEmpty()) {
	String allFailuresMessage =
	refsFailures.stream()
	.map(refPair -> String.format("Failed to fetch ref %s", refPair.compareRef.getName()))
	.collect(Collectors.joining(", "));
	Exception firstFailureException = refsFailures.get(0).exception;

	logger.atSevere().withCause(firstFailureException).log(allFailuresMessage);
	throw new IOException(allFailuresMessage, firstFailureException);
	}

	try (CloseableSet<AutoCloseable> locks = new CloseableSet<>()) {
	final List<RefPair> finalRefsToUpdate = compareAndGetLatestLocalRefs(refsToUpdate, locks);
	delegateUpdate.invoke();
	boolean sharedDbUpdateSucceeded = false;
	try {
	updateSharedRefDb(batchRefUpdate.getCommands().stream(), finalRefsToUpdate);
	sharedDbUpdateSucceeded = true;
	} catch (Exception e) {
	logger.atWarning().withCause(e).log(
	"Batch ref-update failed because of failure during the global refdb update.");
	} finally {
	if (!sharedDbUpdateSucceeded) {
	List<ReceiveCommand> receiveCommands = batchRefUpdate.getCommands();
	logger.atWarning().log(
	"Batch ref-update failed, set all commands Result to LOCK_FAILURE [%d]",
	commands.size());
	rollback(delegateUpdateRollback, finalRefsToUpdate, receiveCommands);
	}
	}
	} catch (OutOfSyncException e) {
	List<ReceiveCommand> receiveCommands = batchRefUpdate.getCommands();
	logger.atWarning().withCause(e).log(
	"Batch ref-update failing because node is out of sync with the shared ref-db. Set all commands Result to LOCK_FAILURE [%d]",
	receiveCommands.size());
	receiveCommands.forEach((command) -> command.setResult(ReceiveCommand.Result.LOCK_FAILURE));
	}
	}

	private void rollback(
	OneParameterVoidFunction<List<ReceiveCommand>> delegateUpdateRollback,
	List<RefPair> refsBeforeUpdate,
	List<ReceiveCommand> receiveCommands)
	throws IOException {
	List<ReceiveCommand> rollbackCommands =
	refsBeforeUpdate.stream()
	.map(
	refBeforeUpdate ->
	new ReceiveCommand(
	refBeforeUpdate.putValue,
	refBeforeUpdate.compareRef.getObjectId(),
	refBeforeUpdate.getName()))
	.collect(Collectors.toList());
	delegateUpdateRollback.invoke(rollbackCommands);
	receiveCommands.forEach(command -> command.setResult(ReceiveCommand.Result.LOCK_FAILURE));
	}

	private void updateSharedRefDb(Stream<ReceiveCommand> commandStream, List<RefPair> refsToUpdate)
	throws IOException {
	if (commandStream.anyMatch(cmd -> cmd.getResult() != ReceiveCommand.Result.OK)) {
	return;
	}

	for (RefPair refPair : refsToUpdate) {
	updateSharedDbOrThrowExceptionFor(refPair);
	}
	}

	private Stream<RefPair> getRefPairs(List<ReceiveCommand> receivedCommands) {
	return receivedCommands.stream().map(this::getRefPairForCommand);
	}

	private RefPair getRefPairForCommand(ReceiveCommand command) {
	try {
	switch (command.getType()) {
	case CREATE:
	return new RefPair(nullRef(command.getRefName()), getNewRef(command));

	case UPDATE:
	case UPDATE_NONFASTFORWARD:
	return new RefPair(getCurrentRef(command.getRefName()), getNewRef(command));

	case DELETE:
	return new RefPair(getCurrentRef(command.getRefName()), ObjectId.zeroId());

	default:
	return new RefPair(
	command.getRef(),
	new IllegalArgumentException("Unsupported command type " + command.getType()));
	}
	} catch (IOException e) {
	return new RefPair(command.getRef(), e);
	}
	}

	private ObjectId getNewRef(ReceiveCommand command) {
	return command.getNewId();
	}

	private List<RefPair> compareAndGetLatestLocalRefs(
	List<RefPair> refsToUpdate, CloseableSet<AutoCloseable> locks) throws IOException {
	List<RefPair> latestRefsToUpdate = new ArrayList<>();
	for (RefPair refPair : refsToUpdate) {
	latestRefsToUpdate.add(compareAndGetLatestLocalRef(refPair, locks));
	}
	return latestRefsToUpdate;
	}

	private static final Ref nullRef(String refName) {
	return new ObjectIdRef.Unpeeled(Ref.Storage.NETWORK, refName, ObjectId.zeroId());
	}
	}