org.apache.aurora.scheduler.storage.durability.DataCompatibilityTest.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.aurora.scheduler.storage.durability.DataCompatibilityTest.java

Source

/**
 * 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 org.apache.aurora.scheduler.storage.durability;

import java.io.File;
import java.io.IOException;
import java.io.UncheckedIOException;
import java.lang.reflect.Method;
import java.lang.reflect.Type;
import java.net.URL;
import java.util.AbstractMap.SimpleImmutableEntry;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import java.util.stream.Stream;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.MapDifference;
import com.google.common.collect.MapDifference.ValueDifference;
import com.google.common.collect.Maps;
import com.google.common.collect.Ordering;
import com.google.common.collect.Streams;
import com.google.common.io.Files;
import com.google.gson.GsonBuilder;
import com.google.gson.JsonParser;
import com.google.inject.AbstractModule;
import com.google.inject.Guice;
import com.google.inject.Injector;

import org.apache.aurora.common.inject.Bindings;
import org.apache.aurora.common.stats.StatsProvider;
import org.apache.aurora.gen.Resource;
import org.apache.aurora.gen.ResourceAggregate;
import org.apache.aurora.gen.storage.Op;
import org.apache.aurora.gen.storage.PruneJobUpdateHistory;
import org.apache.aurora.gen.storage.RemoveJob;
import org.apache.aurora.gen.storage.RemoveJobUpdates;
import org.apache.aurora.gen.storage.RemoveLock;
import org.apache.aurora.gen.storage.RemoveQuota;
import org.apache.aurora.gen.storage.RemoveTasks;
import org.apache.aurora.gen.storage.SaveCronJob;
import org.apache.aurora.gen.storage.SaveFrameworkId;
import org.apache.aurora.gen.storage.SaveHostAttributes;
import org.apache.aurora.gen.storage.SaveJobInstanceUpdateEvent;
import org.apache.aurora.gen.storage.SaveJobUpdate;
import org.apache.aurora.gen.storage.SaveJobUpdateEvent;
import org.apache.aurora.gen.storage.SaveLock;
import org.apache.aurora.gen.storage.SaveQuota;
import org.apache.aurora.gen.storage.SaveTasks;
import org.apache.aurora.scheduler.TierInfo;
import org.apache.aurora.scheduler.TierManager.TierManagerImpl.TierConfig;
import org.apache.aurora.scheduler.TierModule;
import org.apache.aurora.scheduler.app.LifecycleModule;
import org.apache.aurora.scheduler.events.EventSink;
import org.apache.aurora.scheduler.storage.Storage.NonVolatileStorage;
import org.apache.aurora.scheduler.storage.Storage.Volatile;
import org.apache.aurora.scheduler.storage.mem.MemStorageModule;
import org.apache.aurora.scheduler.testing.FakeStatsProvider;
import org.apache.thrift.TDeserializer;
import org.apache.thrift.TException;
import org.apache.thrift.TSerializer;
import org.apache.thrift.TUnion;
import org.apache.thrift.protocol.TJSONProtocol;
import org.junit.Test;

import static com.google.common.base.Charsets.UTF_8;

import static org.apache.aurora.scheduler.storage.durability.Generator.newStruct;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.fail;

public class DataCompatibilityTest {

    private NonVolatileStorage createStorage(Persistence persistence) {
        Injector injector = Guice
                .createInjector(new DurableStorageModule(),
                        new MemStorageModule(Bindings.annotatedKeyFactory(Volatile.class)), new LifecycleModule(),
                        new TierModule(new TierConfig("string-value",
                                ImmutableMap.of("string-value", new TierInfo(false, false)))),
                        new AbstractModule() {
                            @Override
                            protected void configure() {
                                bind(StatsProvider.class).toInstance(new FakeStatsProvider());
                                bind(EventSink.class).toInstance(event -> {
                                });
                                bind(Persistence.class).toInstance(persistence);
                            }
                        });
        return injector.getInstance(NonVolatileStorage.class);
    }

    /**
     * Ops to serve as a reference for the replacement golden values when read compatibility changes.
     */
    private static final List<Op> READ_COMPATIBILITY_OPS = ImmutableList.of(
            Op.pruneJobUpdateHistory(newStruct(PruneJobUpdateHistory.class)),
            Op.removeJob(newStruct(RemoveJob.class)), Op.removeJobUpdate(newStruct(RemoveJobUpdates.class)),
            Op.removeLock(newStruct(RemoveLock.class)), Op.removeQuota(newStruct(RemoveQuota.class)),
            Op.removeTasks(newStruct(RemoveTasks.class)), Op.saveCronJob(newStruct(SaveCronJob.class)),
            Op.saveFrameworkId(newStruct(SaveFrameworkId.class)),
            Op.saveHostAttributes(newStruct(SaveHostAttributes.class)),
            Op.saveJobUpdate(newStruct(SaveJobUpdate.class)),
            Op.saveJobInstanceUpdateEvent(newStruct(SaveJobInstanceUpdateEvent.class)),
            Op.saveJobUpdateEvent(newStruct(SaveJobUpdateEvent.class)), Op.saveLock(newStruct(SaveLock.class)),
            Op.saveQuota(new SaveQuota().setRole("role")
                    .setQuota(new ResourceAggregate().setResources(
                            ImmutableSet.of(Resource.numCpus(2.0), Resource.diskMb(1), Resource.ramMb(1))))),
            Op.saveTasks(newStruct(SaveTasks.class)));

    @Test
    public void testReadCompatibility() {
        // Verifies that storage can recover known-good serialized records.  A failure of this test case
        // indicates that the scheduler can no longer read a record that it was expected to in the past.
        // Golden values in `goldens/read-compatible` preserve serialized records that the scheduler is
        // expected to read.  At the end of a deprecation cycle, these files may need to be updated.
        // Golden file names are prefixed with an ordering ID (e.g. 2-removeJob) to prescribe a recovery
        // order.  This is This is necessary to accommodate Ops with relations (e.g. update events
        // relate to an update).

        // Sanity check that the current read-compatibility values can be replayed.
        NonVolatileStorage storage = createStorage(new TestPersistence(READ_COMPATIBILITY_OPS));
        storage.prepare();
        storage.start(stores -> {
        });
        storage.stop();

        File goldensDir = getGoldensDir("read-compatible");
        List<Op> goldenOps = loadGoldenSchemas(goldensDir).entrySet().stream()
                .sorted(Ordering.natural().onResultOf(entry -> Integer.parseInt(entry.getKey().split("\\-")[0])))
                .map(Entry::getValue).map(DataCompatibilityTest::deserialize).collect(Collectors.toList());

        // Ensure all currently-known Op types are represented in the goldens.
        assertEquals(ImmutableSet.copyOf(Op._Fields.values()),
                goldenOps.stream().map(TUnion::getSetField).collect(Collectors.toSet()));

        // Introduce each op one at a time to pinpoint a specific failed op.
        IntStream.range(1, goldenOps.size()).forEach(i -> {
            NonVolatileStorage store = createStorage(new TestPersistence(goldenOps.subList(0, i)));
            store.prepare();
            try {
                store.start(stores -> {
                });
            } catch (RuntimeException e) {
                Op failedOp = goldenOps.get(i - 1);
                Op currentOp = READ_COMPATIBILITY_OPS.stream()
                        .filter(op -> op.getSetField() == failedOp.getSetField()).findFirst().get();
                StringBuilder error = new StringBuilder().append("**** Storage compatibility change detected ****")
                        .append("\nFailed to recover when introducing ")
                        .append(failedOp.getSetField().getFieldName()).append("\n").append(failedOp)
                        .append("\nIf this is expected, you may delete the associated golden file from ")
                        .append(goldensDir.getPath())
                        .append(",\nor you may replace the file with the latest serialized value:").append("\n")
                        .append(serialize(currentOp));
                fail(error.toString());
            }
            store.stop();
        });
    }

    private static File getGoldensDir(String kind) {
        ClassLoader loader = Thread.currentThread().getContextClassLoader();
        URL url = loader.getResource(
                DataCompatibilityTest.class.getPackage().getName().replaceAll("\\.", "/") + "/goldens/" + kind);
        return new File(url.getPath());
    }

    private static Map<String, String> loadGoldenSchemas(File goldensDir) {
        return Stream.of(goldensDir.listFiles()).collect(Collectors.toMap(File::getName, goldenFile -> {
            try {
                return Files.asCharSource(goldenFile, UTF_8).read();
            } catch (IOException e) {
                throw new UncheckedIOException(e);
            }
        }));
    }

    private static Map<String, String> generateOpSchemas() {
        return Stream.of(Op._Fields.values()).map(field -> {
            Method factory = Stream.of(Op.class.getDeclaredMethods())
                    .filter(method -> method.getName().equals(field.getFieldName())).findFirst().get();

            Class<?> paramType = factory.getParameterTypes()[0];
            Type genericParamType = factory.getGenericParameterTypes()[0];
            try {
                return (Op) factory.invoke(null, Generator.valueFor(paramType, genericParamType));
            } catch (ReflectiveOperationException e) {
                throw new RuntimeException(e);
            }
        }).collect(Collectors.toMap(op -> op.getSetField().getFieldName(), DataCompatibilityTest::serialize));
    }

    @Test
    public void testWriteFormatUnchanged() {
        // Attempts to flag any changes in the storage format.  While thorough, this check is not
        // complete.  It attempts to capture the entire schema by synthesizing a fully-populated
        // instance of each Op type.  For TUnions, the struct generator picks an arbitrary field to set,
        // meaning that it will only see one of the multiple possible schemas for any given TUnion.
        // These generated structs effectively give a view of the struct schema, which is compared to
        // golden files in `goldens/current`.

        Map<String, String> schemasByName = generateOpSchemas();
        File goldensDir = getGoldensDir("current");
        Map<String, String> goldensByName = loadGoldenSchemas(goldensDir);

        MapDifference<String, String> difference = Maps.difference(goldensByName, schemasByName);
        if (difference.areEqual()) {
            return;
        }

        StringBuilder error = new StringBuilder();
        StringBuilder remedy = new StringBuilder();

        Set<String> removedOps = difference.entriesOnlyOnLeft().keySet();
        if (!removedOps.isEmpty()) {
            error.append("Removal of storage Op(s): ").append(removedOps)
                    .append("\nOps may only be removed after a release that")
                    .append("\n  * formally deprecates the Op in release notes")
                    .append("\n  * performs a no-op read of the Op type")
                    .append("\n  * included warning logging when the Op was read")
                    .append("\n  * ensures the Op is removed from storage")
                    .append("\n\nHowever, you should also consider leaving the Op indefinitely and removing")
                    .append("\nall fields as a safer alternative.");

            remedy.append("deleting the files ")
                    .append(removedOps.stream().map(removed -> new File(goldensDir, removed).getAbsolutePath())
                            .collect(Collectors.joining(", ")));
        }

        String goldenChangeInstructions = Streams
                .concat(difference.entriesOnlyOnRight().entrySet().stream(),
                        difference.entriesDiffering().entrySet().stream().map(
                                entry -> new SimpleImmutableEntry<>(entry.getKey(), entry.getValue().rightValue())))
                .map(entry -> new StringBuilder().append("\n")
                        .append(new File(goldensDir, entry.getKey()).getPath()).append(":").append("\n")
                        .append(entry.getValue()).toString())
                .collect(Collectors.joining("\n"));

        Set<String> addedOps = difference.entriesOnlyOnRight().keySet();
        if (!addedOps.isEmpty()) {
            error.append("Addition of storage Op(s): ").append(addedOps).append("\nOps may only be introduced")
                    .append("\n  a.) in a release that supports reading but not writing the Op")
                    .append("\n  b.) in a release that writes the Op only with an operator-controlled flag");

            remedy.append("creating the following files").append(goldenChangeInstructions);
        }

        Map<String, ValueDifference<String>> modified = difference.entriesDiffering();
        if (!modified.isEmpty()) {
            error.append("Schema changes to Op(s): " + modified.keySet())
                    .append("\nThis check detects that changes occured, not how the schema changed.")
                    .append("\nSome guidelines for evolving schemas:")
                    .append("\n  * Introducing fields: you must handle reading records that do not")
                    .append("\n    yet have the field set.  This can be done with a backfill routine during")
                    .append("\n    storage recovery if a field is required in some parts of the code")
                    .append("\n  * Removing fields: must only be done after a release in which the field")
                    .append("\n    is unused and announced as deprecated")
                    .append("\n  * Changed fields: the type or thrift field ID of a field must never change");

            remedy.append("changing the following files").append(goldenChangeInstructions);
        }

        fail(new StringBuilder().append("**** Storage compatibility change detected ****").append("\n")
                .append(error).append("\n\nIf the necessary compatibility procedures have been performed,")
                .append("\nyou may clear this check by ").append(remedy).toString());
    }

    private static class TestPersistence implements Persistence {
        private final List<Op> ops;

        TestPersistence(List<Op> ops) {
            this.ops = ops;
        }

        @Override
        public void prepare() {
            // No-op.
        }

        @Override
        public Stream<Edit> recover() {
            return ops.stream().map(Edit::op);
        }

        @Override
        public void persist(Stream<Op> records) {
            // no-op.
        }
    }

    private static String serialize(Op op) {
        try {
            String unformattedJson = new String(new TSerializer(new TJSONProtocol.Factory()).serialize(op), UTF_8);

            // Pretty print the json for easier review of diffs.
            return new GsonBuilder().setPrettyPrinting().create().toJson(new JsonParser().parse(unformattedJson))
                    + "\n";
        } catch (TException e) {
            throw new RuntimeException(e);
        }
    }

    private static Op deserialize(String serializedOp) {
        try {
            Op op = new Op();

            String nonPrettyJson = new GsonBuilder().create().toJson(new JsonParser().parse(serializedOp));

            new TDeserializer(new TJSONProtocol.Factory()).deserialize(op, nonPrettyJson.getBytes(UTF_8));
            return op;
        } catch (TException e) {
            throw new RuntimeException(e);
        }
    }
}