io.pravega.client.stream.mock.MockController.java Source code

Java tutorial

Introduction

Here is the source code for io.pravega.client.stream.mock.MockController.java

Source

/**
 * Copyright (c) 2017 Dell Inc., or its subsidiaries. All Rights Reserved.
 *
 * 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
 */
package io.pravega.client.stream.mock;

import com.google.common.base.Preconditions;
import io.pravega.client.netty.impl.ClientConnection;
import io.pravega.client.netty.impl.ConnectionFactory;
import io.pravega.client.segment.impl.Segment;
import io.pravega.client.stream.ScalingPolicy;
import io.pravega.client.stream.Stream;
import io.pravega.client.stream.StreamConfiguration;
import io.pravega.client.stream.Transaction;
import io.pravega.client.stream.TxnFailedException;
import io.pravega.client.stream.impl.CancellableRequest;
import io.pravega.client.stream.impl.ConnectionClosedException;
import io.pravega.client.stream.impl.Controller;
import io.pravega.client.stream.impl.StreamCut;
import io.pravega.client.stream.impl.StreamImpl;
import io.pravega.client.stream.impl.StreamSegments;
import io.pravega.client.stream.impl.StreamSegmentsWithPredecessors;
import io.pravega.client.stream.impl.TxnSegments;
import io.pravega.common.concurrent.FutureHelpers;
import io.pravega.shared.protocol.netty.FailingReplyProcessor;
import io.pravega.shared.protocol.netty.PravegaNodeUri;
import io.pravega.shared.protocol.netty.ReplyProcessor;
import io.pravega.shared.protocol.netty.WireCommand;
import io.pravega.shared.protocol.netty.WireCommands;
import io.pravega.shared.protocol.netty.WireCommands.AbortTransaction;
import io.pravega.shared.protocol.netty.WireCommands.CommitTransaction;
import io.pravega.shared.protocol.netty.WireCommands.CreateSegment;
import io.pravega.shared.protocol.netty.WireCommands.CreateTransaction;
import io.pravega.shared.protocol.netty.WireCommands.DeleteSegment;
import io.pravega.shared.protocol.netty.WireCommands.TransactionAborted;
import io.pravega.shared.protocol.netty.WireCommands.TransactionCommitted;
import io.pravega.shared.protocol.netty.WireCommands.TransactionCreated;
import io.pravega.shared.protocol.netty.WireCommands.WrongHost;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeMap;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Supplier;
import java.util.stream.Collectors;
import javax.annotation.concurrent.GuardedBy;
import lombok.AllArgsConstructor;
import lombok.Synchronized;
import org.apache.commons.lang.NotImplementedException;

import static io.pravega.common.concurrent.FutureHelpers.getAndHandleExceptions;

@AllArgsConstructor
public class MockController implements Controller {

    private final String endpoint;
    private final int port;
    private final ConnectionFactory connectionFactory;
    @GuardedBy("$lock")
    private final Map<String, Set<Stream>> createdScopes = new HashMap<>();
    @GuardedBy("$lock")
    private final Map<Stream, StreamConfiguration> createdStreams = new HashMap<>();
    private final Supplier<Long> idGenerator = new AtomicLong(0)::incrementAndGet;

    @Override
    @Synchronized
    public CompletableFuture<Boolean> createScope(final String scopeName) {
        if (createdScopes.get(scopeName) != null) {
            return CompletableFuture.completedFuture(false);
        }
        createdScopes.put(scopeName, new HashSet<>());
        return CompletableFuture.completedFuture(true);
    }

    @Override
    @Synchronized
    public CompletableFuture<Boolean> deleteScope(String scopeName) {
        if (createdScopes.get(scopeName) == null) {
            return CompletableFuture.completedFuture(false);
        }

        if (!createdScopes.get(scopeName).isEmpty()) {
            return FutureHelpers.failedFuture(new IllegalStateException("Scope is not empty."));
        }

        createdScopes.remove(scopeName);
        return CompletableFuture.completedFuture(true);
    }

    @Override
    @Synchronized
    public CompletableFuture<Boolean> createStream(StreamConfiguration streamConfig) {
        Stream stream = new StreamImpl(streamConfig.getScope(), streamConfig.getStreamName());
        if (createdStreams.get(stream) != null) {
            return CompletableFuture.completedFuture(false);
        }

        if (createdScopes.get(streamConfig.getScope()) == null) {
            return FutureHelpers.failedFuture(new IllegalArgumentException("Scope does not exit."));
        }

        createdStreams.put(stream, streamConfig);
        createdScopes.get(streamConfig.getScope()).add(stream);
        for (Segment segment : getSegmentsForStream(stream)) {
            createSegment(segment.getScopedName(), new PravegaNodeUri(endpoint, port));
        }
        return CompletableFuture.completedFuture(true);
    }

    @Synchronized
    List<Segment> getSegmentsForStream(Stream stream) {
        StreamConfiguration config = createdStreams.get(stream);
        Preconditions.checkArgument(config != null, "Stream must be created first");
        ScalingPolicy scalingPolicy = config.getScalingPolicy();
        if (scalingPolicy.getType() != ScalingPolicy.Type.FIXED_NUM_SEGMENTS) {
            throw new IllegalArgumentException("Dynamic scaling not supported with a mock controller");
        }
        List<Segment> result = new ArrayList<>(scalingPolicy.getMinNumSegments());
        for (int i = 0; i < scalingPolicy.getMinNumSegments(); i++) {
            result.add(new Segment(config.getScope(), config.getStreamName(), i));
        }
        return result;
    }

    @Override
    public CompletableFuture<Boolean> updateStream(StreamConfiguration streamConfig) {
        throw new NotImplementedException();
    }

    @Override
    public CompletableFuture<Boolean> startScale(Stream stream, List<Integer> sealedSegments,
            Map<Double, Double> newKeyRanges) {
        throw new NotImplementedException();
    }

    @Override
    public CancellableRequest<Boolean> scaleStream(Stream stream, List<Integer> sealedSegments,
            Map<Double, Double> newKeyRanges, ScheduledExecutorService executor) {
        throw new NotImplementedException();
    }

    @Override
    public CompletableFuture<Boolean> checkScaleStatus(Stream stream, int epoch) {
        throw new NotImplementedException();
    }

    @Override
    public CompletableFuture<Boolean> sealStream(String scope, String streamName) {
        throw new NotImplementedException();
    }

    @Override
    @Synchronized
    public CompletableFuture<Boolean> deleteStream(String scope, String streamName) {
        Stream stream = new StreamImpl(scope, streamName);
        if (createdStreams.get(stream) == null) {
            return CompletableFuture.completedFuture(false);
        }
        for (Segment segment : getSegmentsForStream(stream)) {
            deleteSegment(segment.getScopedName(), new PravegaNodeUri(endpoint, port));
        }
        createdStreams.remove(stream);
        createdScopes.get(scope).remove(stream);
        return CompletableFuture.completedFuture(true);
    }

    private boolean createSegment(String name, PravegaNodeUri uri) {
        CompletableFuture<Boolean> result = new CompletableFuture<>();
        FailingReplyProcessor replyProcessor = new FailingReplyProcessor() {

            @Override
            public void connectionDropped() {
                result.completeExceptionally(new ConnectionClosedException());
            }

            @Override
            public void wrongHost(WireCommands.WrongHost wrongHost) {
                result.completeExceptionally(new NotImplementedException());
            }

            @Override
            public void segmentAlreadyExists(WireCommands.SegmentAlreadyExists segmentAlreadyExists) {
                result.complete(false);
            }

            @Override
            public void segmentCreated(WireCommands.SegmentCreated segmentCreated) {
                result.complete(true);
            }

            @Override
            public void processingFailure(Exception error) {
                result.completeExceptionally(error);
            }
        };
        CreateSegment command = new WireCommands.CreateSegment(idGenerator.get(), name,
                WireCommands.CreateSegment.NO_SCALE, 0);
        sendRequestOverNewConnection(command, replyProcessor, result);
        return getAndHandleExceptions(result, RuntimeException::new);
    }

    private boolean deleteSegment(String name, PravegaNodeUri uri) {
        CompletableFuture<Boolean> result = new CompletableFuture<>();
        FailingReplyProcessor replyProcessor = new FailingReplyProcessor() {

            @Override
            public void connectionDropped() {
                result.completeExceptionally(new ConnectionClosedException());
            }

            @Override
            public void wrongHost(WireCommands.WrongHost wrongHost) {
                result.completeExceptionally(new NotImplementedException());
            }

            @Override
            public void segmentDeleted(WireCommands.SegmentDeleted segmentDeleted) {
                result.complete(true);
            }

            @Override
            public void noSuchSegment(WireCommands.NoSuchSegment noSuchSegment) {
                result.complete(false);
            }

            @Override
            public void processingFailure(Exception error) {
                result.completeExceptionally(error);
            }
        };
        DeleteSegment command = new WireCommands.DeleteSegment(idGenerator.get(), name);
        sendRequestOverNewConnection(command, replyProcessor, result);
        return getAndHandleExceptions(result, RuntimeException::new);
    }

    @Override
    public CompletableFuture<StreamSegments> getCurrentSegments(String scope, String stream) {
        return CompletableFuture.completedFuture(getCurrentSegments(new StreamImpl(scope, stream)));
    }

    private StreamSegments getCurrentSegments(Stream stream) {
        List<Segment> segmentsInStream = getSegmentsForStream(stream);
        TreeMap<Double, Segment> segments = new TreeMap<>();
        double increment = 1.0 / segmentsInStream.size();
        for (int i = 0; i < segmentsInStream.size(); i++) {
            segments.put((i + 1) * increment, new Segment(stream.getScope(), stream.getStreamName(), i));
        }
        return new StreamSegments(segments);
    }

    @Override
    public CompletableFuture<Void> commitTransaction(Stream stream, UUID txId) {
        List<CompletableFuture<Void>> futures = new ArrayList<>();
        for (Segment segment : getSegmentsForStream(stream)) {
            futures.add(commitTxSegment(txId, segment));
        }
        return FutureHelpers.allOf(futures);
    }

    private CompletableFuture<Void> commitTxSegment(UUID txId, Segment segment) {
        CompletableFuture<Void> result = new CompletableFuture<>();
        FailingReplyProcessor replyProcessor = new FailingReplyProcessor() {

            @Override
            public void connectionDropped() {
                result.completeExceptionally(new ConnectionClosedException());
            }

            @Override
            public void wrongHost(WrongHost wrongHost) {
                result.completeExceptionally(new NotImplementedException());
            }

            @Override
            public void transactionCommitted(TransactionCommitted transactionCommitted) {
                result.complete(null);
            }

            @Override
            public void transactionAborted(TransactionAborted transactionAborted) {
                result.completeExceptionally(new TxnFailedException("Transaction already aborted."));
            }

            @Override
            public void processingFailure(Exception error) {
                result.completeExceptionally(error);
            }
        };
        sendRequestOverNewConnection(new CommitTransaction(idGenerator.get(), segment.getScopedName(), txId),
                replyProcessor, result);
        return result;
    }

    @Override
    public CompletableFuture<Void> abortTransaction(Stream stream, UUID txId) {
        List<CompletableFuture<Void>> futures = new ArrayList<>();
        for (Segment segment : getSegmentsForStream(stream)) {
            futures.add(abortTxSegment(txId, segment));
        }
        return FutureHelpers.allOf(futures);
    }

    private CompletableFuture<Void> abortTxSegment(UUID txId, Segment segment) {
        CompletableFuture<Void> result = new CompletableFuture<>();
        FailingReplyProcessor replyProcessor = new FailingReplyProcessor() {

            @Override
            public void connectionDropped() {
                result.completeExceptionally(new ConnectionClosedException());
            }

            @Override
            public void wrongHost(WrongHost wrongHost) {
                result.completeExceptionally(new NotImplementedException());
            }

            @Override
            public void transactionCommitted(TransactionCommitted transactionCommitted) {
                result.completeExceptionally(new RuntimeException("Transaction already committed."));
            }

            @Override
            public void transactionAborted(TransactionAborted transactionAborted) {
                result.complete(null);
            }

            @Override
            public void processingFailure(Exception error) {
                result.completeExceptionally(error);
            }
        };
        sendRequestOverNewConnection(new AbortTransaction(idGenerator.get(), segment.getScopedName(), txId),
                replyProcessor, result);
        return result;
    }

    @Override
    public CompletableFuture<Transaction.Status> checkTransactionStatus(Stream stream, UUID txId) {
        throw new NotImplementedException();
    }

    @Override
    public CompletableFuture<TxnSegments> createTransaction(final Stream stream, final long lease,
            final long maxExecutionTime, final long scaleGracePeriod) {
        UUID txId = UUID.randomUUID();
        List<CompletableFuture<Void>> futures = new ArrayList<>();
        StreamSegments currentSegments = getCurrentSegments(stream);
        for (Segment segment : currentSegments.getSegments()) {
            futures.add(createSegmentTx(txId, segment));
        }
        return FutureHelpers.allOf(futures).thenApply(v -> new TxnSegments(currentSegments, txId));
    }

    private CompletableFuture<Void> createSegmentTx(UUID txId, Segment segment) {
        CompletableFuture<Void> result = new CompletableFuture<>();
        FailingReplyProcessor replyProcessor = new FailingReplyProcessor() {

            @Override
            public void connectionDropped() {
                result.completeExceptionally(new ConnectionClosedException());
            }

            @Override
            public void wrongHost(WrongHost wrongHost) {
                result.completeExceptionally(new NotImplementedException());
            }

            @Override
            public void transactionCreated(TransactionCreated transactionCreated) {
                result.complete(null);
            }

            @Override
            public void processingFailure(Exception error) {
                result.completeExceptionally(error);
            }
        };
        sendRequestOverNewConnection(new CreateTransaction(idGenerator.get(), segment.getScopedName(), txId),
                replyProcessor, result);
        return result;
    }

    @Override
    public CompletableFuture<Void> pingTransaction(Stream stream, UUID txId, long lease) {
        throw new NotImplementedException();
    }

    @Override
    public CompletableFuture<Map<Segment, Long>> getSegmentsAtTime(Stream stream, long timestamp) {
        return CompletableFuture
                .completedFuture(getSegmentsForStream(stream).stream().collect(Collectors.toMap(s -> s, s -> 0L)));
    }

    @Override
    public CompletableFuture<StreamSegmentsWithPredecessors> getSuccessors(Segment segment) {
        return CompletableFuture.completedFuture(new StreamSegmentsWithPredecessors(Collections.emptyMap()));
    }

    @Override
    public CompletableFuture<Set<Segment>> getSuccessors(StreamCut from) {
        throw new NotImplementedException();
    }

    @Override
    public CompletableFuture<PravegaNodeUri> getEndpointForSegment(String qualifiedSegmentName) {
        return CompletableFuture.completedFuture(new PravegaNodeUri(endpoint, port));
    }

    private <T> void sendRequestOverNewConnection(WireCommand request, ReplyProcessor replyProcessor,
            CompletableFuture<T> resultFuture) {
        ClientConnection connection = getAndHandleExceptions(
                connectionFactory.establishConnection(new PravegaNodeUri(endpoint, port), replyProcessor),
                RuntimeException::new);
        resultFuture.whenComplete((result, e) -> {
            connection.close();
        });
        try {
            connection.send(request);
        } catch (Exception e) {
            resultFuture.completeExceptionally(e);
        }
    }

    @Override
    public CompletableFuture<Boolean> isSegmentOpen(Segment segment) {
        return CompletableFuture.completedFuture(true);
    }

}