org.apache.cassandra.streaming.StreamCoordinator Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of cassandra-all Show documentation
The Apache Cassandra Project develops a highly scalable second-generation distributed database, bringing together Dynamo's fully distributed design and Bigtable's ColumnFamily-based data model.
There is a newer version: 5.0.2
Show newest version
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.cassandra.streaming;

import java.net.InetSocketAddress;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

import com.google.common.annotations.VisibleForTesting;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.cassandra.locator.InetAddressAndPort;
import org.apache.cassandra.utils.TimeUUID;

import static org.apache.cassandra.net.MessagingService.current_version;


/**
 * {@link StreamCoordinator} is a helper class that abstracts away maintaining multiple
 * StreamSession and ProgressInfo instances per peer.
 *
 * This class coordinates multiple SessionStreams per peer in both the outgoing StreamPlan context and on the
 * inbound StreamResultFuture context.
 */
public class StreamCoordinator
{
    private static final Logger logger = LoggerFactory.getLogger(StreamCoordinator.class);

    private final boolean connectSequentially;

    private final Map peerSessions = new ConcurrentHashMap<>();
    private final StreamOperation streamOperation;
    private final int connectionsPerHost;
    private final boolean follower;
    private StreamingChannel.Factory factory;
    private Iterator sessionsToConnect = null;
    private final TimeUUID pendingRepair;
    private final PreviewKind previewKind;

    public StreamCoordinator(StreamOperation streamOperation, int connectionsPerHost, StreamingChannel.Factory factory,
                             boolean follower, boolean connectSequentially, TimeUUID pendingRepair, PreviewKind previewKind)
    {
        this.streamOperation = streamOperation;
        this.connectionsPerHost = connectionsPerHost;
        this.factory = factory;
        this.follower = follower;
        this.connectSequentially = connectSequentially;
        this.pendingRepair = pendingRepair;
        this.previewKind = previewKind;
    }

    public void setConnectionFactory(StreamingChannel.Factory factory)
    {
        this.factory = factory;
    }

    /**
     * @return true if any stream session is active
     */
    public synchronized boolean hasActiveSessions()
    {
        for (HostStreamingData data : peerSessions.values())
        {
            if (data.hasActiveSessions())
                return true;
        }
        return false;
    }

    public synchronized Collection getAllStreamSessions()
    {
        Collection results = new ArrayList<>();
        for (HostStreamingData data : peerSessions.values())
        {
            results.addAll(data.getAllStreamSessions());
        }
        return results;
    }

    public boolean isFollower()
    {
        return follower;
    }

    public void connect(StreamResultFuture future)
    {
        if (this.connectSequentially)
            connectSequentially(future);
        else
            connectAllStreamSessions();
    }

    private void connectAllStreamSessions()
    {
        for (HostStreamingData data : peerSessions.values())
            data.connectAllStreamSessions();
    }

    private void connectSequentially(StreamResultFuture future)
    {
        sessionsToConnect = getAllStreamSessions().iterator();
        future.addEventListener(new StreamEventHandler()
        {
            public void handleStreamEvent(StreamEvent event)
            {
                if (event.eventType == StreamEvent.Type.STREAM_PREPARED || event.eventType == StreamEvent.Type.STREAM_COMPLETE)
                    connectNext();
            }

            public void onSuccess(StreamState result)
            {

            }

            public void onFailure(Throwable t)
            {

            }
        });
        connectNext();
    }

    private void connectNext()
    {
        if (sessionsToConnect == null)
            return;

        if (sessionsToConnect.hasNext())
        {
            StreamSession next = sessionsToConnect.next();
            if (logger.isDebugEnabled())
                logger.debug("Connecting next session {} with {}.", next.planId(), next.peer.toString());
            startSession(next);
        }
        else
            logger.debug("Finished connecting all sessions");
    }

    public synchronized Set getPeers()
    {
        return new HashSet<>(peerSessions.keySet());
    }

    public synchronized StreamSession getOrCreateOutboundSession(InetAddressAndPort peer)
    {
        return getOrCreateHostData(peer).getOrCreateOutboundSession(peer);
    }

    public synchronized StreamSession getOrCreateInboundSession(InetAddressAndPort from, StreamingChannel channel, int messagingVersion, int id)
    {
        return getOrCreateHostData(from).getOrCreateInboundSession(from, channel, messagingVersion, id);
    }

    public StreamSession getSessionById(InetAddressAndPort peer, int id)
    {
        return getHostData(peer).getSessionById(id);
    }

    public synchronized void updateProgress(ProgressInfo info)
    {
        getHostData(info.peer).updateProgress(info);
    }

    public synchronized void addSessionInfo(SessionInfo session)
    {
        HostStreamingData data = getOrCreateHostData(session.peer);
        data.addSessionInfo(session);
    }

    public synchronized Set getAllSessionInfo()
    {
        Set result = new HashSet<>();
        for (HostStreamingData data : peerSessions.values())
        {
            result.addAll(data.getAllSessionInfo());
        }
        return result;
    }

    public synchronized void transferStreams(InetAddressAndPort to, Collection streams)
    {
        HostStreamingData sessionList = getOrCreateHostData(to);

        if (connectionsPerHost > 1)
        {
            List> buckets = bucketStreams(streams);

            for (Collection bucket : buckets)
            {
                StreamSession session = sessionList.getOrCreateOutboundSession(to);
                session.addTransferStreams(bucket);
            }
        }
        else
        {
            StreamSession session = sessionList.getOrCreateOutboundSession(to);
            session.addTransferStreams(streams);
        }
    }

    private List> bucketStreams(Collection streams)
    {
        // There's no point in divvying things up into more buckets than we have sstableDetails
        int targetSlices = Math.min(streams.size(), connectionsPerHost);
        int step = Math.round((float) streams.size() / (float) targetSlices);
        int index = 0;

        List> result = new ArrayList<>();
        List slice = null;

        for (OutgoingStream stream: streams)
        {
            if (index % step == 0)
            {
                slice = new ArrayList<>();
                result.add(slice);
            }
            slice.add(stream);
            ++index;
        }
        return result;
    }

    private HostStreamingData getHostData(InetAddressAndPort peer)
    {
        HostStreamingData data = peerSessions.get(peer);

        if (data == null)
            throw new IllegalArgumentException("Unknown peer requested: " + peer);
        return data;
    }

    private HostStreamingData getOrCreateHostData(InetSocketAddress peer)
    {
        HostStreamingData data = peerSessions.get(peer);
        if (data == null)
        {
            data = new HostStreamingData();
            peerSessions.put(peer, data);
        }
        return data;
    }

    public TimeUUID getPendingRepair()
    {
        return pendingRepair;
    }

    private void startSession(StreamSession session)
    {
        session.start();
        logger.info("[Stream #{}, ID#{}] Beginning stream session with {}", session.planId(), session.sessionIndex(), session.peer);
    }

    private class HostStreamingData
    {
        private final Map streamSessions = new HashMap<>();
        private final Map sessionInfos = new HashMap<>();

        private int lastReturned = -1;

        public boolean hasActiveSessions()
        {
            for (StreamSession session : streamSessions.values())
            {
                if (!session.state().isFinalState())
                    return true;
            }
            return false;
        }

        public StreamSession getOrCreateOutboundSession(InetAddressAndPort peer)
        {
            // create
            if (streamSessions.size() < connectionsPerHost)
            {
                StreamSession session = new StreamSession(streamOperation, peer, factory, null, current_version, isFollower(), streamSessions.size(),
                                                          pendingRepair, previewKind);
                streamSessions.put(++lastReturned, session);
                sessionInfos.put(lastReturned, session.getSessionInfo());
                return session;
            }
            // get
            else
            {
                if (lastReturned >= streamSessions.size() - 1)
                    lastReturned = 0;

                return streamSessions.get(lastReturned++);
            }
        }

        public void connectAllStreamSessions()
        {
            for (StreamSession session : streamSessions.values())
            {
                startSession(session);
            }
        }

        public Collection getAllStreamSessions()
        {
            return Collections.unmodifiableCollection(streamSessions.values());
        }

        public StreamSession getOrCreateInboundSession(InetAddressAndPort from, StreamingChannel channel, int messagingVersion, int id)
        {
            StreamSession session = streamSessions.get(id);
            if (session == null)
            {
                session = new StreamSession(streamOperation, from, factory, channel, messagingVersion, isFollower(), id, pendingRepair, previewKind);
                streamSessions.put(id, session);
                sessionInfos.put(id, session.getSessionInfo());
            }
            return session;
        }

        public StreamSession getSessionById(int id)
        {
            return streamSessions.get(id);
        }

        public void updateProgress(ProgressInfo info)
        {
            sessionInfos.get(info.sessionIndex).updateProgress(info);
        }

        public void addSessionInfo(SessionInfo info)
        {
            sessionInfos.put(info.sessionIndex, info);
        }

        public Collection getAllSessionInfo()
        {
            return sessionInfos.values();
        }

        @VisibleForTesting
        public void shutdown()
        {
            streamSessions.values().forEach(ss -> ss.sessionFailed());
        }
    }
}