com.hubspot.singularity.data.zkmigrations.PendingRequestDataMigration Maven / Gradle / Ivy

Go to download
package com.hubspot.singularity.data.zkmigrations;

import com.google.inject.Inject;
import com.hubspot.mesos.JavaUtils;
import com.hubspot.singularity.SingularityDeployKey;
import com.hubspot.singularity.SingularityPendingRequest;
import com.hubspot.singularity.SingularityPendingRequest.PendingType;
import com.hubspot.singularity.data.RequestManager;
import com.hubspot.singularity.data.transcoders.Transcoder;
import java.util.List;
import org.apache.curator.framework.CuratorFramework;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class PendingRequestDataMigration extends ZkDataMigration {
  private static final Logger LOG = LoggerFactory.getLogger(
    PendingRequestDataMigration.class
  );

  private final RequestManager requestManager;
  private final CuratorFramework curator;
  private final Transcoder requestTranscoder;

  @Inject
  public PendingRequestDataMigration(
    RequestManager requestManager,
    CuratorFramework curator,
    Transcoder requestTranscoder
  ) {
    super(10);
    this.requestManager = requestManager;
    this.curator = curator;
    this.requestTranscoder = requestTranscoder;
  }

  @Override
  public void applyMigration() {
    String basePath = "/requests/pending";
    LOG.info("Starting migration to re-write pending request paths");
    long start = System.currentTimeMillis();
    int rewrittenPaths = 0;

    try {
      if (curator.checkExists().forPath(basePath) == null) {
        return;
      }
    } catch (Exception exn) {
      LOG.error("Could not check existence of pending request path", exn);
      throw new RuntimeException(exn);
    }

    try {
      List childPaths = curator.getChildren().forPath(basePath);

      for (String childPath : childPaths) {
        SingularityPendingRequest pendingRequest = requestTranscoder.fromBytes(
          curator.getData().forPath(String.format("%s/%s", basePath, childPath))
        );
        if (pendingRequest.getPendingType() == PendingType.IMMEDIATE) {
          String rewrittenPath = new SingularityDeployKey(
            pendingRequest.getRequestId(),
            pendingRequest.getDeployId()
          )
          .getId();
          LOG.warn(
            "Rewriting path {} to {}",
            childPath,
            String.format("%s%s", rewrittenPath, pendingRequest.getTimestamp())
          );
          requestManager.addToPendingQueue(pendingRequest);
          curator.delete().forPath(String.format("%s/%s", basePath, childPath));
          rewrittenPaths += 1;
        } else {
          LOG.warn("Not rewriting path {}, already correct", childPath);
        }
      }
    } catch (Exception exn) {
      LOG.error("Connection to Zookeeper failed while running migration", exn);
      throw new RuntimeException(exn);
    }

    LOG.warn(
      "Applied PendingRequestDataMigration to {} requests in {}",
      rewrittenPaths,
      JavaUtils.duration(start)
    );
  }
}