UdpSocketReceiver.h

Go to the documentation of this file.

// Copyright Epic Games, Inc. All Rights Reserved.
 
#pragma once
 
#include "CoreMinimal.h"
#include "HAL/Runnable.h"
#include "HAL/RunnableThread.h"
#include "Misc/SingleThreadRunnable.h"
#include "Serialization/ArrayReader.h"
#include "Sockets.h"
#include "SocketSubsystem.h"
 
#include "Interfaces/IPv4/IPv4Endpoint.h"
#include "IPAddress.h"
 
#include <atomic>
 
typedef TSharedPtr<FArrayReader, ESPMode::ThreadSafe> FArrayReaderPtr;
 
DECLARE_DELEGATE_TwoParams(FOnSocketDataReceived, const FArrayReaderPtr&, const FIPv4Endpoint&);
 
 
class FUdpSocketReceiver
    : public FRunnable
    , private FSingleThreadRunnable
{
public:
 
    FUdpSocketReceiver(FSocket* InSocket, const FTimespan& InWaitTime, const TCHAR* InThreadName)
        : Socket(InSocket)
        , bStopping(false)
        , Thread(nullptr)
        , ThreadName(InThreadName)
        , WaitTime(InWaitTime)
    {
        check(Socket != nullptr);
        check(Socket->GetSocketType() == SOCKTYPE_Datagram);
 
        SocketSubsystem = ISocketSubsystem::Get(PLATFORM_SOCKETSUBSYSTEM);
    }
 
    virtual ~FUdpSocketReceiver()
    {
        if (Thread != nullptr)
        {
            Thread->Kill(true);
            delete Thread;
        }
    }
 
public:
 
    void SetMaxReadBufferSize(uint32 InMaxReadBufferSize)
    {
        MaxReadBufferSize = InMaxReadBufferSize;
    }
 
    void SetThreadStackSize(uint32 InTheadStackSize)
    {
        ThreadStackSize = InTheadStackSize;
    }
 
    void Start()
    {
        Thread = FRunnableThread::Create(this, *ThreadName, ThreadStackSize, TPri_AboveNormal, FPlatformAffinity::GetPoolThreadMask());
    }
 
    FOnSocketDataReceived& OnDataReceived()
    {
        check(Thread == nullptr);
        return DataReceivedDelegate;
    }
 
public:
 
    //~ FRunnable interface
 
    virtual FSingleThreadRunnable* GetSingleThreadInterface() override
    {
        return this;
    }
 
    virtual bool Init() override
    {
        return true;
    }
 
    virtual uint32 Run() override
    {
        while (!bStopping.load(std::memory_order_acquire))
        {
            Update(WaitTime);
        }
 
        return 0;
    }
 
    virtual void Stop() override
    {
        bStopping.store(true, std::memory_order_release);
    }
 
    virtual void Exit() override { }
 
protected:
 
    void Update(const FTimespan& SocketWaitTime)
    {
        if (!Socket->Wait(ESocketWaitConditions::WaitForRead, SocketWaitTime))
        {
            return;
        }
 
        TRACE_CPUPROFILER_EVENT_SCOPE(FUdpSocketReceiver_Update);
        TSharedRef<FInternetAddr> Sender = SocketSubsystem->CreateInternetAddr();
        uint32 Size;
 
        while (Socket->HasPendingData(Size))
        {
            FArrayReaderPtr Reader = MakeShared<FArrayReader, ESPMode::ThreadSafe>(true);
            Reader->SetNumUninitialized(FMath::Min(Size, MaxReadBufferSize));
 
            int32 Read = 0;
 
            if (Socket->RecvFrom(Reader->GetData(), Reader->Num(), Read, *Sender))
            {
                ensure((uint32)Read <= MaxReadBufferSize);
                Reader->RemoveAt(Read, Reader->Num() - Read, EAllowShrinking::No);
                DataReceivedDelegate.ExecuteIfBound(Reader, FIPv4Endpoint(Sender));
            }
        }
    }
 
protected:
 
    //~ FSingleThreadRunnable interface
 
    virtual void Tick() override
    {
        Update(FTimespan::Zero());
    }
 
private:
 
    FSocket* Socket;
 
    ISocketSubsystem* SocketSubsystem;
 
    std::atomic<bool> bStopping;
 
    FRunnableThread* Thread;
 
    FString ThreadName;
 
    FTimespan WaitTime;
 
    uint32 MaxReadBufferSize = 65507u;
 
    uint32 ThreadStackSize = 128 * 1024;
 
private:
 
    FOnSocketDataReceived DataReceivedDelegate;
};