Examples with SetMultimap com.google.common.collect.SetMultimap used on opensource projects

Example 6 with SetMultimap

use of com.google.common.collect.SetMultimap in project presto by prestodb.

the class TestPrestoSparkSourceDistributionSplitAssigner method assertSplitAssignment.

private static void assertSplitAssignment(boolean autoTuneEnabled, DataSize maxSplitsDataSizePerSparkPartition, int initialPartitionCount, int minSparkInputPartitionCountForAutoTune, int maxSparkInputPartitionCountForAutoTune, List<Long> splitSizes, Map<Integer, List<Long>> expectedAssignment) {
    // assign splits in one shot
    {
        PrestoSparkSplitAssigner assigner = new PrestoSparkSourceDistributionSplitAssigner(new PlanNodeId("test"), createSplitSource(splitSizes), Integer.MAX_VALUE, maxSplitsDataSizePerSparkPartition.toBytes(), initialPartitionCount, autoTuneEnabled, minSparkInputPartitionCountForAutoTune, maxSparkInputPartitionCountForAutoTune);
        Optional<SetMultimap<Integer, ScheduledSplit>> actualAssignment = assigner.getNextBatch();
        if (!splitSizes.isEmpty()) {
            assertThat(actualAssignment).isPresent();
            assertAssignedSplits(actualAssignment.get(), expectedAssignment);
        } else {
            assertThat(actualAssignment).isNotPresent();
        }
    }
    // assign splits iteratively
    for (int splitBatchSize = 1; splitBatchSize < splitSizes.size(); splitBatchSize *= 2) {
        HashMultimap<Integer, ScheduledSplit> actualAssignment = HashMultimap.create();
        // sort splits to make assignment match the assignment done in one shot
        List<Long> sortedSplits = new ArrayList<>(splitSizes);
        sortedSplits.sort(Comparator.<Long>naturalOrder().reversed());
        PrestoSparkSplitAssigner assigner = new PrestoSparkSourceDistributionSplitAssigner(new PlanNodeId("test"), createSplitSource(sortedSplits), splitBatchSize, maxSplitsDataSizePerSparkPartition.toBytes(), initialPartitionCount, autoTuneEnabled, minSparkInputPartitionCountForAutoTune, maxSparkInputPartitionCountForAutoTune);
        while (true) {
            Optional<SetMultimap<Integer, ScheduledSplit>> assignment = assigner.getNextBatch();
            if (!assignment.isPresent()) {
                break;
            }
            actualAssignment.putAll(assignment.get());
        }
        assertAssignedSplits(actualAssignment, expectedAssignment);
    }
}

Example 7 with SetMultimap

use of com.google.common.collect.SetMultimap in project presto by prestodb.

the class DiscoveryNodeManager method refreshNodesInternal.

private synchronized void refreshNodesInternal() {
    // This is currently a blacklist.
    // TODO: make it a whitelist (a failure-detecting service selector) and maybe build in support for injecting this in airlift
    Set<ServiceDescriptor> services = serviceSelector.selectAllServices().stream().filter(service -> !failureDetector.getFailed().contains(service)).filter(service -> !nodeStatusService.isPresent() || nodeStatusService.get().isAllowed(service.getLocation()) || isCoordinator(service) || isResourceManager(service)).collect(toImmutableSet());
    ImmutableSet.Builder<InternalNode> activeNodesBuilder = ImmutableSortedSet.orderedBy(comparing(InternalNode::getNodeIdentifier));
    ImmutableSet.Builder<InternalNode> inactiveNodesBuilder = ImmutableSet.builder();
    ImmutableSet.Builder<InternalNode> shuttingDownNodesBuilder = ImmutableSet.builder();
    ImmutableSet.Builder<InternalNode> coordinatorsBuilder = ImmutableSet.builder();
    ImmutableSet.Builder<InternalNode> resourceManagersBuilder = ImmutableSet.builder();
    ImmutableSetMultimap.Builder<ConnectorId, InternalNode> byConnectorIdBuilder = ImmutableSetMultimap.builder();
    Map<String, InternalNode> nodes = new HashMap<>();
    SetMultimap<String, ConnectorId> connectorIdsByNodeId = HashMultimap.create();
    // For a given connectorId, sort the nodes based on their nodeIdentifier
    byConnectorIdBuilder.orderValuesBy(comparing(InternalNode::getNodeIdentifier));
    if (isMemoizeDeadNodesEnabled && this.nodes != null) {
        nodes.putAll(this.nodes);
    }
    if (isMemoizeDeadNodesEnabled && this.connectorIdsByNodeId != null) {
        connectorIdsByNodeId.putAll(this.connectorIdsByNodeId);
    }
    for (ServiceDescriptor service : services) {
        URI uri = getHttpUri(service, httpsRequired);
        OptionalInt thriftPort = getThriftServerPort(service);
        NodeVersion nodeVersion = getNodeVersion(service);
        // Currently, a node may have the roles of both a coordinator and a worker.  In the future, a resource manager may also
        // take the form of a coordinator, hence these flags are not exclusive.
        boolean coordinator = isCoordinator(service);
        boolean resourceManager = isResourceManager(service);
        if (uri != null && nodeVersion != null) {
            InternalNode node = new InternalNode(service.getNodeId(), uri, thriftPort, nodeVersion, coordinator, resourceManager, ALIVE);
            NodeState nodeState = getNodeState(node);
            switch(nodeState) {
                case ACTIVE:
                    activeNodesBuilder.add(node);
                    if (coordinator) {
                        coordinatorsBuilder.add(node);
                    }
                    if (resourceManager) {
                        resourceManagersBuilder.add(node);
                    }
                    nodes.put(node.getNodeIdentifier(), node);
                    // record available active nodes organized by connector id
                    String connectorIds = service.getProperties().get("connectorIds");
                    if (connectorIds != null) {
                        connectorIds = connectorIds.toLowerCase(ENGLISH);
                        for (String id : CONNECTOR_ID_SPLITTER.split(connectorIds)) {
                            ConnectorId connectorId = new ConnectorId(id);
                            byConnectorIdBuilder.put(connectorId, node);
                            connectorIdsByNodeId.put(node.getNodeIdentifier(), connectorId);
                        }
                    }
                    // always add system connector
                    byConnectorIdBuilder.put(new ConnectorId(GlobalSystemConnector.NAME), node);
                    break;
                case INACTIVE:
                    inactiveNodesBuilder.add(node);
                    break;
                case SHUTTING_DOWN:
                    shuttingDownNodesBuilder.add(node);
                    break;
                default:
                    log.error("Unknown state %s for node %s", nodeState, node);
            }
        }
    }
    if (allNodes != null) {
        // log node that are no longer active (but not shutting down)
        SetView<InternalNode> missingNodes = difference(allNodes.getActiveNodes(), Sets.union(activeNodesBuilder.build(), shuttingDownNodesBuilder.build()));
        for (InternalNode missingNode : missingNodes) {
            log.info("Previously active node is missing: %s (last seen at %s)", missingNode.getNodeIdentifier(), missingNode.getHost());
        }
    }
    // nodes by connector id changes anytime a node adds or removes a connector (note: this is not part of the listener system)
    activeNodesByConnectorId = byConnectorIdBuilder.build();
    if (isMemoizeDeadNodesEnabled) {
        SetView<String> deadNodeIds = difference(nodes.keySet(), activeNodesBuilder.build().stream().map(InternalNode::getNodeIdentifier).collect(toImmutableSet()));
        for (String nodeId : deadNodeIds) {
            InternalNode deadNode = nodes.get(nodeId);
            Set<ConnectorId> deadNodeConnectorIds = connectorIdsByNodeId.get(nodeId);
            for (ConnectorId id : deadNodeConnectorIds) {
                byConnectorIdBuilder.put(id, new InternalNode(deadNode.getNodeIdentifier(), deadNode.getInternalUri(), deadNode.getThriftPort(), deadNode.getNodeVersion(), deadNode.isCoordinator(), deadNode.isResourceManager(), DEAD));
            }
        }
    }
    this.nodes = ImmutableMap.copyOf(nodes);
    this.nodesByConnectorId = byConnectorIdBuilder.build();
    this.connectorIdsByNodeId = ImmutableSetMultimap.copyOf(connectorIdsByNodeId);
    AllNodes allNodes = new AllNodes(activeNodesBuilder.build(), inactiveNodesBuilder.build(), shuttingDownNodesBuilder.build(), coordinatorsBuilder.build(), resourceManagersBuilder.build());
    // only update if all nodes actually changed (note: this does not include the connectors registered with the nodes)
    if (!allNodes.equals(this.allNodes)) {
        // assign allNodes to a local variable for use in the callback below
        this.allNodes = allNodes;
        coordinators = coordinatorsBuilder.build();
        resourceManagers = resourceManagersBuilder.build();
        // notify listeners
        List<Consumer<AllNodes>> listeners = ImmutableList.copyOf(this.listeners);
        nodeStateEventExecutor.submit(() -> listeners.forEach(listener -> listener.accept(allNodes)));
    }
}

Example 8 with SetMultimap

use of com.google.common.collect.SetMultimap in project cassandra by apache.

the class StartupClusterConnectivityChecker method execute.

/**
 * @param peers The currently known peers in the cluster; argument is not modified.
 * @param getDatacenterSource A function for mapping peers to their datacenter.
 * @return true if the requested percentage of peers are marked ALIVE in gossip and have their connections opened;
 * else false.
 */
public boolean execute(Set<InetAddressAndPort> peers, Function<InetAddressAndPort, String> getDatacenterSource) {
    if (peers == null || this.timeoutNanos < 0)
        return true;
    // make a copy of the set, to avoid mucking with the input (in case it's a sensitive collection)
    peers = new HashSet<>(peers);
    InetAddressAndPort localAddress = FBUtilities.getBroadcastAddressAndPort();
    String localDc = getDatacenterSource.apply(localAddress);
    peers.remove(localAddress);
    if (peers.isEmpty())
        return true;
    // make a copy of the datacenter mapping (in case gossip updates happen during this method or some such)
    Map<InetAddressAndPort, String> peerToDatacenter = new HashMap<>();
    SetMultimap<String, InetAddressAndPort> datacenterToPeers = HashMultimap.create();
    for (InetAddressAndPort peer : peers) {
        String datacenter = getDatacenterSource.apply(peer);
        peerToDatacenter.put(peer, datacenter);
        datacenterToPeers.put(datacenter, peer);
    }
    // on the remaining local datacenter.
    if (!blockForRemoteDcs) {
        datacenterToPeers.keySet().retainAll(Collections.singleton(localDc));
        logger.info("Blocking coordination until only a single peer is DOWN in the local datacenter, timeout={}s", TimeUnit.NANOSECONDS.toSeconds(timeoutNanos));
    } else {
        logger.info("Blocking coordination until only a single peer is DOWN in each datacenter, timeout={}s", TimeUnit.NANOSECONDS.toSeconds(timeoutNanos));
    }
    // The threshold is 3 because for each peer we want to have 3 acks,
    // one for small message connection, one for large message connnection and one for alive event from gossip.
    AckMap acks = new AckMap(3, peers);
    Map<String, CountDownLatch> dcToRemainingPeers = new HashMap<>(datacenterToPeers.size());
    for (String datacenter : datacenterToPeers.keys()) {
        dcToRemainingPeers.put(datacenter, newCountDownLatch(Math.max(datacenterToPeers.get(datacenter).size() - 1, 0)));
    }
    long startNanos = nanoTime();
    // set up a listener to react to new nodes becoming alive (in gossip), and account for all the nodes that are already alive
    Set<InetAddressAndPort> alivePeers = Collections.newSetFromMap(new ConcurrentHashMap<>());
    AliveListener listener = new AliveListener(alivePeers, dcToRemainingPeers, acks, peerToDatacenter::get);
    Gossiper.instance.register(listener);
    // send out a ping message to open up the non-gossip connections to all peers. Note that this sends the
    // ping messages to _all_ peers, not just the ones we block for in dcToRemainingPeers.
    sendPingMessages(peers, dcToRemainingPeers, acks, peerToDatacenter::get);
    for (InetAddressAndPort peer : peers) {
        if (Gossiper.instance.isAlive(peer) && alivePeers.add(peer) && acks.incrementAndCheck(peer)) {
            String datacenter = peerToDatacenter.get(peer);
            // We have to check because we might only have the local DC in the map
            if (dcToRemainingPeers.containsKey(datacenter))
                dcToRemainingPeers.get(datacenter).decrement();
        }
    }
    boolean succeeded = true;
    for (CountDownLatch countDownLatch : dcToRemainingPeers.values()) {
        long remainingNanos = Math.max(1, timeoutNanos - (nanoTime() - startNanos));
        // noinspection UnstableApiUsage
        succeeded &= countDownLatch.awaitUninterruptibly(remainingNanos, TimeUnit.NANOSECONDS);
    }
    Gossiper.instance.unregister(listener);
    if (succeeded) {
        logger.info("Ensured sufficient healthy connections with {} after {} milliseconds", dcToRemainingPeers.keySet(), TimeUnit.NANOSECONDS.toMillis(nanoTime() - startNanos));
    } else {
        // dc -> missing peer host addresses
        Map<String, List<String>> peersDown = acks.getMissingPeers().stream().collect(groupingBy(peer -> {
            String dc = peerToDatacenter.get(peer);
            if (dc != null)
                return dc;
            return StringUtils.defaultString(getDatacenterSource.apply(peer), "unknown");
        }, mapping(InetAddressAndPort::getHostAddressAndPort, toList())));
        logger.warn("Timed out after {} milliseconds, was waiting for remaining peers to connect: {}", TimeUnit.NANOSECONDS.toMillis(nanoTime() - startNanos), peersDown);
    }
    return succeeded;
}

Example 9 with SetMultimap

use of com.google.common.collect.SetMultimap in project tiger by google.

the class Tiger3ProcessorForComponent method addChildrenAndRemoveSub.

private void addChildrenAndRemoveSub(Set<TypeElement> allEitherComponents, Map<TypeElement, TypeElement> componentToParentMap) {
    SetMultimap eitherComponentToChildrenMap = HashMultimap.create();
    for (Map.Entry<TypeElement, TypeElement> i : componentToParentMap.entrySet()) {
        eitherComponentToChildrenMap.put(i.getValue(), i.getKey());
    }
    Set<TypeElement> toAdd = new HashSet<>();
    for (TypeElement i : allEitherComponents) {
        toAdd.addAll(eitherComponentToChildrenMap.get(i));
    }
    Set<TypeElement> toRemove = new HashSet<>();
    for (TypeElement i : allEitherComponents) {
        if (utils.isSubcomponent(i)) {
            toRemove.add(i);
        }
    }
    for (TypeElement i : toRemove) {
        allEitherComponents.remove(i);
    }
    for (TypeElement i : toAdd) {
        allEitherComponents.add(i);
    }
}

Example 10 with SetMultimap

use of com.google.common.collect.SetMultimap in project tiger by google.

the class Tiger4ProcessorForComponent method addChildrenAndRemoveSub.

private void addChildrenAndRemoveSub(Set<TypeElement> allEitherComponents, Map<TypeElement, TypeElement> componentToParentMap) {
    SetMultimap eitherComponentToChildrenMap = HashMultimap.create();
    for (Map.Entry<TypeElement, TypeElement> i : componentToParentMap.entrySet()) {
        eitherComponentToChildrenMap.put(i.getValue(), i.getKey());
    }
    Set<TypeElement> toAdd = new HashSet<>();
    for (TypeElement i : allEitherComponents) {
        toAdd.addAll(eitherComponentToChildrenMap.get(i));
    }
    Set<TypeElement> toRemove = new HashSet<>();
    for (TypeElement i : allEitherComponents) {
        if (utils.isSubcomponent(i)) {
            toRemove.add(i);
        }
    }
    for (TypeElement i : toRemove) {
        allEitherComponents.remove(i);
    }
    for (TypeElement i : toAdd) {
        allEitherComponents.add(i);
    }
}