Skip to content
Permalink
Browse files
Cassandra diff should accept a provided job_id for retrying diffs
Patch by Jyothsna Konisa; reviewed by Dinesh Joshi, Yifan Cai for CASSANDRA-16968
  • Loading branch information
jyothsnakonisa authored and yifan-c committed Sep 22, 2021
1 parent e3161b6 commit 75b70c35eea537d038c9f72632eec3d141b05379
Showing 4 changed files with 89 additions and 10 deletions.
@@ -67,5 +67,12 @@
<artifactId>junit</artifactId>
</dependency>

<dependency>
<groupId>org.mockito</groupId>
<artifactId>mockito-core</artifactId>
<version>3.5.10</version>
<scope>test</scope>
</dependency>

</dependencies>
</project>
@@ -122,9 +122,12 @@ public void run(JobConfiguration configuration, JavaSparkContext sc) {
ClusterProvider metadataProvider = ClusterProvider.getProvider(configuration.clusterConfig("metadata"), "metadata");
JobMetadataDb.JobLifeCycle job = null;
UUID jobId = null;
try (Cluster metadataCluster = metadataProvider.getCluster();
Session metadataSession = metadataCluster.connect()) {
Cluster metadataCluster = null;
Session metadataSession = null;

try {
metadataCluster = metadataProvider.getCluster();
metadataSession = metadataCluster.connect();
RetryStrategyProvider retryStrategyProvider = RetryStrategyProvider.create(configuration.retryOptions());
MetadataKeyspaceOptions metadataOptions = configuration.metadataOptions();
JobMetadataDb.Schema.maybeInitialize(metadataSession, metadataOptions, retryStrategyProvider);
@@ -197,18 +200,32 @@ public void run(JobConfiguration configuration, JavaSparkContext sc) {
Differ.shutdown();
JobMetadataDb.ProgressTracker.resetStatements();
}
if (metadataCluster != null) {
metadataCluster.close();
}
if (metadataSession != null) {
metadataSession.close();
}

}
}

private static Params getJobParams(JobMetadataDb.JobLifeCycle job, JobConfiguration conf, List<KeyspaceTablePair> keyspaceTables) {
@VisibleForTesting
static Params getJobParams(JobMetadataDb.JobLifeCycle job, JobConfiguration conf, List<KeyspaceTablePair> keyspaceTables) {
if (conf.jobId().isPresent()) {
return job.getJobParams(conf.jobId().get());
} else {
return new Params(UUID.randomUUID(),
keyspaceTables,
conf.buckets(),
conf.splits());
final Params jobParams = job.getJobParams(conf.jobId().get());
if(jobParams != null) {
// When job_id is passed as a config property for the first time, we will not have metadata associated
// with job_id in metadata table. we should return jobParams from the table only when jobParams is not null
// Otherwise return new jobParams with provided job_id
return jobParams;
}
}
final UUID jobId = conf.jobId().isPresent() ? conf.jobId().get() : UUID.randomUUID();
return new Params(jobId,
keyspaceTables,
conf.buckets(),
conf.splits());
}

private static List<Split> getSplits(JobConfiguration config, TokenHelper tokenHelper) {
@@ -369,7 +369,7 @@ public void initializeJob(DiffJob.Params params,
metadataKeyspace, Schema.RUNNING_JOBS),
params.jobId);
if (!rs.one().getBool("[applied]")) {
logger.info("Aborting due to inability to mark job as running. " +
logger.info("Could not mark job as running. " +
"Did a previous run of job id {} fail non-gracefully?",
params.jobId);
throw new RuntimeException("Unable to mark job running, aborting");
@@ -20,11 +20,18 @@
package org.apache.cassandra.diff;

import java.math.BigInteger;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.UUID;

import org.junit.Test;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

public class DiffJobTest
{
@@ -39,6 +46,37 @@ public void testSplitsRandom()
splitTestHelper(TokenHelper.forPartitioner("RandomPartitioner"));
}

@Test
public void testGetJobParamsWithJobIdProvidedShouldReturnNonNullConFigParams() {
final MockConfig mockConfig = new MockConfig();
final JobMetadataDb.JobLifeCycle mockJob = mock(JobMetadataDb.JobLifeCycle.class);
final List<KeyspaceTablePair> keyspaceTablePairs = new ArrayList<>();
final DiffJob.Params params = DiffJob.getJobParams(mockJob, mockConfig, keyspaceTablePairs);
assertNotNull(params);
}

@Test
public void testGetJobParamsDuringRetryShouldReturnPreviousParams() {
final MockConfig mockConfig = new MockConfig();
final JobMetadataDb.JobLifeCycle mockJob = mock(JobMetadataDb.JobLifeCycle.class);
final DiffJob.Params mockParams = mock(DiffJob.Params.class);
when(mockJob.getJobParams(any())).thenAnswer(invocationOnMock -> mockParams);
final List<KeyspaceTablePair> keyspaceTablePairs = new ArrayList<>();
final DiffJob.Params params = DiffJob.getJobParams(mockJob, mockConfig, keyspaceTablePairs);
assertEquals(params, mockParams);
}

@Test
public void testGetJobParamsWithNoJobId() {
final MockConfig mockConfig = mock(MockConfig.class);
when(mockConfig.jobId()).thenReturn(Optional.empty());

final JobMetadataDb.JobLifeCycle mockJob = mock(JobMetadataDb.JobLifeCycle.class);
final List<KeyspaceTablePair> keyspaceTablePairs = new ArrayList<>();
final DiffJob.Params params = DiffJob.getJobParams(mockJob, mockConfig, keyspaceTablePairs);
assertNotNull(params.jobId);
}

private void splitTestHelper(TokenHelper tokens)
{
List<DiffJob.Split> splits = DiffJob.calculateSplits(50, 1, tokens);
@@ -54,4 +92,21 @@ private void splitTestHelper(TokenHelper tokens)
for (int i = 0; i < splits.size(); i++)
assertEquals(i, splits.get(i).splitNumber);
}

private class MockConfig extends AbstractMockJobConfiguration {
@Override
public int splits() {
return 2;
}

@Override
public int buckets() {
return 2;
}

@Override
public Optional<UUID> jobId() {
return Optional.of(UUID.randomUUID());
}
}
}

0 comments on commit 75b70c3

Please sign in to comment.