MultichannelTcpSocket.h

Go to the documentation of this file.

// Copyright Epic Games, Inc. All Rights Reserved.
 
#pragma once
 
#include "CoreMinimal.h"
#include "Serialization/BufferArchive.h"
#include "Misc/ScopedEvent.h"
#include "Misc/ScopeLock.h"
#include "Serialization/MemoryReader.h"
#include "MultichannelTcpReceiver.h"
#include "MultichannelTcpSender.h"
 
class FSocket;
 
class FMultichannelTcpSocket
{
    enum
    {
        ControlChannel = 0
    };
 
public:
 
    FMultichannelTcpSocket( FSocket* InSocket, uint64 InBandwidthLatencyProduct )
        : BandwidthLatencyProduct(InBandwidthLatencyProduct)
        , RemoteReceiverBytesReceived(0)
        , Receiver(InSocket, FOnMultichannelTcpReceive::CreateRaw(this, &FMultichannelTcpSocket::HandleReceiverReceive))
        , Sender(InSocket, FOnMultichannelTcpOkToSend::CreateRaw(this, &FMultichannelTcpSocket::HandleSenderOkToSend))
        , Socket(InSocket)
    { }
 
public:
 
    int32 BlockingReceive( uint8* Data, int32 Count, uint32 Channel )
    {
        check((Channel != ControlChannel) && Data); // Channel 0 is reserved, must want data
        
        if (Count > 0)
        {
            for (int32 Attempt = 0; Attempt < 2; Attempt++)
            {
                FScopedEvent* LocalEventToRestart = NULL;
                {
                    FScopeLock ScopeLock(&ReceiveBuffersCriticalSection);
 
                    FReceiveBuffer* ChannelBuffer = ReceiveBuffers.FindRef(Channel);
 
                    if (!ChannelBuffer)
                    {
                        ChannelBuffer = new FReceiveBuffer();
 
                        ReceiveBuffers.Add(Channel, ChannelBuffer);
                    }
 
                    // Would be bad to have multiple listeners
                    check(!ChannelBuffer->EventToResumeWhenDataIsReady && !ChannelBuffer->BytesRequiredToResume);
                
                    if (ChannelBuffer->Buffer.Num() >= Count)
                    {
                        FMemory::Memcpy(Data, ChannelBuffer->Buffer.GetData(), Count);
 
                        if (ChannelBuffer->Buffer.Num() == Count)
                        {
                            ReceiveBuffers.Remove(Channel);
                        }
                        else
                        {
                            ChannelBuffer->Buffer.RemoveAt(0, Count);
                        }
 
                        return Count;
                    }
 
                    if (!Attempt) // if someone woke us up with insufficient data, we simply error out
                    {
                        LocalEventToRestart = new FScopedEvent();
 
                        ChannelBuffer->EventToResumeWhenDataIsReady = LocalEventToRestart;
                        ChannelBuffer->BytesRequiredToResume = Count;
                    }
                }
 
                delete LocalEventToRestart; // wait here for sufficient data
            }
        }
 
        return 0;
    }
 
    int32 DataAvailable( uint32 Channel )
    {
        FScopeLock ScopeLock(&ReceiveBuffersCriticalSection);
 
        FReceiveBuffer* ChannelBuffer = ReceiveBuffers.FindRef(Channel);
 
        return ChannelBuffer ? ChannelBuffer->Buffer.Num() : 0;
    }
 
    int32 PollingReceive( uint8* Data, int32 MaxCount, uint32 Channel )
    {
        check(Channel != ControlChannel && Data && MaxCount > 0); // Channel 0 is reserved, must want data
 
        FScopeLock ScopeLock(&ReceiveBuffersCriticalSection);
 
        FReceiveBuffer* ChannelBuffer = ReceiveBuffers.FindRef(Channel);
        
        if (!ChannelBuffer)
        {
            return 0;
        }
 
        check(!ChannelBuffer->EventToResumeWhenDataIsReady && !ChannelBuffer->BytesRequiredToResume && ChannelBuffer->Buffer.Num()); // would be bad to have multiple listeners
        
        int32 Count = FMath::Min<int32>(ChannelBuffer->Buffer.Num(), MaxCount);
 
        FMemory::Memcpy(Data, ChannelBuffer->Buffer.GetData(), Count);
        
        if (ChannelBuffer->Buffer.Num() == Count)
        {
            ReceiveBuffers.Remove(Channel);
        }
        else
        {
            ChannelBuffer->Buffer.RemoveAt(0, Count);
        }
 
        return Count;
    }
 
    void Send( const uint8* Data, int32 Count, uint32 Channel )
    {
        check((Channel != ControlChannel) && Data && Count); // Channel 0 is reserved, must have data
        Sender.Send(Data, Count, Channel);
    }
 
private:
 
    struct FReceiveBuffer
    {
        TArray<uint8> Buffer;
 
        int32 BytesRequiredToResume;
 
        FScopedEvent* EventToResumeWhenDataIsReady;
 
        FReceiveBuffer( )
            : BytesRequiredToResume(0)
            , EventToResumeWhenDataIsReady(NULL)
        { }
    };
 
    // Callback for receiving data from the receiver thread.
    void HandleReceiverReceive( const TArray<uint8>& Payload, uint32 Channel, bool bForceByteswapping )
    {
        if (Channel == ControlChannel)
        {
            FMemoryReader Ar(Payload);
            
            Ar.SetByteSwapping(bForceByteswapping);
            
            uint64 RemoteBytesReceived;
            Ar << RemoteBytesReceived;
            
            RemoteReceiverBytesReceived = RemoteBytesReceived;
            
            Sender.AttemptResumeSending();
            
            return;
        }
 
        // process the actual payload
        {
            FScopeLock ScopeLock(&ReceiveBuffersCriticalSection);
            FReceiveBuffer* ChannelBuffer = ReceiveBuffers.FindRef(Channel);
            
            if (!ChannelBuffer)
            {
                ChannelBuffer = new FReceiveBuffer();
                ReceiveBuffers.Add(Channel, ChannelBuffer);
            }
            
            ChannelBuffer->Buffer += Payload;
            
            if (ChannelBuffer->Buffer.Num() >= ChannelBuffer->BytesRequiredToResume && ChannelBuffer->EventToResumeWhenDataIsReady)
            {
                ChannelBuffer->EventToResumeWhenDataIsReady->Trigger();
                ChannelBuffer->EventToResumeWhenDataIsReady = NULL;
                ChannelBuffer->BytesRequiredToResume = 0;
            }
        }
 
        // send a control channel message back
        {
            FBufferArchive Ar;
            uint64 BytesReceived = Receiver.GetBytesReceived();
            Ar << BytesReceived;
            Sender.Send(Ar.GetData(), Ar.Num(), ControlChannel);
        }
    }
 
    bool HandleSenderOkToSend( int32 PayloadSize, uint32 Channel )
    {
        if (Channel == ControlChannel)
        {
            return true; // always ok to send on the control channel
        }
 
        bool Result = (Sender.GetBytesSent() + PayloadSize) < (RemoteReceiverBytesReceived + BandwidthLatencyProduct);
        
        return Result;
    }
 
private:
 
    uint64 BandwidthLatencyProduct;
 
    TMap<uint32, FReceiveBuffer*> ReceiveBuffers;
 
    FCriticalSection ReceiveBuffersCriticalSection;
 
    int64 RemoteReceiverBytesReceived;
 
    FMultichannelTcpReceiver Receiver;
 
    FMultichannelTcpSender Sender;
 
    FSocket* Socket;
};