io.datakernel.datagraph.graph.DataGraph.java Source code

Introduction

Here is the source code for io.datakernel.datagraph.graph.DataGraph.java

Source

/*
 * Copyright (C) 2015 SoftIndex LLC.
 *
 * 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 io.datakernel.datagraph.graph;

import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.ListMultimap;
import com.google.common.reflect.TypeToken;
import io.datakernel.datagraph.node.Node;
import io.datakernel.datagraph.server.DatagraphSerialization;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

/**
 * Represents a graph of partitions, nodes and streams in datagraph system.
 */
@SuppressWarnings("unchecked")
public class DataGraph {
    private final DatagraphSerialization serialization;
    private final List<Partition> availablePartitions;
    private final Map<Node, Partition> nodePartitions = new LinkedHashMap<>();
    private final Map<StreamId, Node> streams = new LinkedHashMap<>();

    public DataGraph(DatagraphSerialization serialization, List<Partition> availablePartitions) {
        this.serialization = serialization;
        this.availablePartitions = availablePartitions;
    }

    public List<Partition> getAvailablePartitions() {
        return availablePartitions;
    }

    public Partition getPartition(Node node) {
        return nodePartitions.get(node);
    }

    public Partition getPartition(StreamId streamId) {
        return getPartition(streams.get(streamId));
    }

    private ListMultimap<Partition, Node> getNodesByPartition() {
        ListMultimap<Partition, Node> listMultimap = ArrayListMultimap.create();
        for (Map.Entry<Node, Partition> entry : nodePartitions.entrySet()) {
            listMultimap.put(entry.getValue(), entry.getKey());
        }
        return listMultimap;
    }

    /**
     * Executes the defined operations on all partitions.
     */
    public void execute() {
        ListMultimap<Partition, Node> map = getNodesByPartition();
        for (Partition partition : map.keySet()) {
            List<Node> nodes = map.get(partition);
            partition.execute(nodes);
        }
    }

    @Override
    public String toString() {
        StringBuilder sb = new StringBuilder();
        ListMultimap<Partition, Node> map = getNodesByPartition();
        for (Partition partition : map.keySet()) {
            sb.append("--- " + partition + "\n\n");
            List<Node> nodes = map.get(partition);
            String str = serialization.gson.toJson(nodes, new TypeToken<List<Node>>() {
            }.getType());
            sb.append(str).append("\n\n");
        }
        return sb.toString();
    }

    public void addNode(Partition partition, Node node) {
        nodePartitions.put(node, partition);
        for (StreamId streamId : node.getOutputs()) {
            streams.put(streamId, node);
        }
    }

    public void addNodeStream(Node node, StreamId streamId) {
        streams.put(streamId, node);
    }

    public List<Partition> getPartitions(List<? extends StreamId> channels) {
        List<Partition> partitions = new ArrayList<>();
        for (StreamId streamId : channels) {
            Partition partition = getPartition(streamId);
            partitions.add(partition);
        }
        return partitions;
    }

    public DatagraphSerialization getSerialization() {
        return serialization;
    }
}