subscription_thread_dispatcher.h

/*
 * Copyright 2025 LiveKit
 *
 * 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.
 */
 
#ifndef LIVEKIT_SUBSCRIPTION_THREAD_DISPATCHER_H
#define LIVEKIT_SUBSCRIPTION_THREAD_DISPATCHER_H
 
#include "livekit/audio_stream.h"
#include "livekit/video_stream.h"
 
#include <cstdint>
#include <functional>
#include <memory>
#include <mutex>
#include <optional>
#include <string>
#include <thread>
#include <unordered_map>
#include <vector>
 
namespace livekit {
 
class AudioFrame;
class DataTrackStream;
class RemoteDataTrack;
class Track;
class VideoFrame;
 
using AudioFrameCallback = std::function<void(const AudioFrame &)>;
 
using VideoFrameCallback =
    std::function<void(const VideoFrame &frame, std::int64_t timestamp_us)>;
 
using DataFrameCallback =
    std::function<void(const std::vector<std::uint8_t> &payload,
                       std::optional<std::uint64_t> user_timestamp)>;
 
using DataFrameCallbackId = std::uint64_t;
 
class SubscriptionThreadDispatcher {
public:
  SubscriptionThreadDispatcher();
 
  ~SubscriptionThreadDispatcher();
 
  void setOnAudioFrameCallback(const std::string &participant_identity,
                               TrackSource source, AudioFrameCallback callback,
                               AudioStream::Options opts = {});
 
  void setOnAudioFrameCallback(const std::string &participant_identity,
                               const std::string &track_name,
                               AudioFrameCallback callback,
                               AudioStream::Options opts = {});
 
  void setOnVideoFrameCallback(const std::string &participant_identity,
                               TrackSource source, VideoFrameCallback callback,
                               VideoStream::Options opts = {});
 
  void setOnVideoFrameCallback(const std::string &participant_identity,
                               const std::string &track_name,
                               VideoFrameCallback callback,
                               VideoStream::Options opts = {});
 
  void clearOnAudioFrameCallback(const std::string &participant_identity,
                                 TrackSource source);
 
  void clearOnAudioFrameCallback(const std::string &participant_identity,
                                 const std::string &track_name);
 
  void clearOnVideoFrameCallback(const std::string &participant_identity,
                                 TrackSource source);
 
  void clearOnVideoFrameCallback(const std::string &participant_identity,
                                 const std::string &track_name);
 
  void handleTrackSubscribed(const std::string &participant_identity,
                             TrackSource source, const std::string &track_name,
                             const std::shared_ptr<Track> &track);
 
  void handleTrackUnsubscribed(const std::string &participant_identity,
                               TrackSource source,
                               const std::string &track_name);
 
  // ---------------------------------------------------------------
  // Data track callbacks
  // ---------------------------------------------------------------
 
  DataFrameCallbackId
  addOnDataFrameCallback(const std::string &participant_identity,
                         const std::string &track_name,
                         DataFrameCallback callback);
 
  void removeOnDataFrameCallback(DataFrameCallbackId id);
 
  void handleDataTrackPublished(const std::shared_ptr<RemoteDataTrack> &track);
 
  void handleDataTrackUnpublished(const std::string &sid);
 
  void stopAll();
 
private:
  friend class SubscriptionThreadDispatcherTest;
 
  struct CallbackKey {
    std::string participant_identity;
    TrackSource source;
    std::string track_name;
 
    bool operator==(const CallbackKey &o) const {
      return participant_identity == o.participant_identity &&
             source == o.source && track_name == o.track_name;
    }
  };
 
  struct CallbackKeyHash {
    std::size_t operator()(const CallbackKey &k) const {
      auto h1 = std::hash<std::string>{}(k.participant_identity);
      auto h2 = std::hash<int>{}(static_cast<int>(k.source));
      auto h3 = std::hash<std::string>{}(k.track_name);
      return h1 ^ (h2 << 1) ^ (h3 << 2);
    }
  };
 
  struct ActiveReader {
    std::shared_ptr<AudioStream> audio_stream;
    std::shared_ptr<VideoStream> video_stream;
    std::thread thread;
  };
 
  struct DataCallbackKey {
    std::string participant_identity;
    std::string track_name;
 
    bool operator==(const DataCallbackKey &o) const {
      return participant_identity == o.participant_identity &&
             track_name == o.track_name;
    }
  };
 
  struct DataCallbackKeyHash {
    std::size_t operator()(const DataCallbackKey &k) const {
      auto h1 = std::hash<std::string>{}(k.participant_identity);
      auto h2 = std::hash<std::string>{}(k.track_name);
      return h1 ^ (h2 << 1);
    }
  };
 
  struct RegisteredDataCallback {
    DataCallbackKey key;
    DataFrameCallback callback;
  };
 
  struct ActiveDataReader {
    std::shared_ptr<RemoteDataTrack> remote_track;
    std::mutex sub_mutex;
    std::shared_ptr<DataTrackStream> stream; // guarded by sub_mutex
    std::thread thread;
  };
 
  struct RegisteredAudioCallback {
    AudioFrameCallback callback;
    AudioStream::Options options;
  };
 
  struct RegisteredVideoCallback {
    VideoFrameCallback callback;
    VideoStream::Options options;
  };
 
  std::thread extractReaderThreadLocked(const CallbackKey &key);
 
  std::thread startReaderLocked(const CallbackKey &key,
                                const std::shared_ptr<Track> &track);
 
  std::thread startAudioReaderLocked(const CallbackKey &key,
                                     const std::shared_ptr<Track> &track,
                                     AudioFrameCallback cb,
                                     const AudioStream::Options &opts);
 
  std::thread startVideoReaderLocked(const CallbackKey &key,
                                     const std::shared_ptr<Track> &track,
                                     VideoFrameCallback cb,
                                     const VideoStream::Options &opts);
 
  std::thread extractDataReaderThreadLocked(DataFrameCallbackId id);
 
  std::thread extractDataReaderThreadLocked(const DataCallbackKey &key);
 
  std::thread
  startDataReaderLocked(DataFrameCallbackId id, const DataCallbackKey &key,
                        const std::shared_ptr<RemoteDataTrack> &track,
                        DataFrameCallback cb);
 
  mutable std::mutex lock_;
 
  std::unordered_map<CallbackKey, RegisteredAudioCallback, CallbackKeyHash>
      audio_callbacks_;
 
  std::unordered_map<CallbackKey, RegisteredVideoCallback, CallbackKeyHash>
      video_callbacks_;
 
  std::unordered_map<CallbackKey, ActiveReader, CallbackKeyHash>
      active_readers_;
 
  DataFrameCallbackId next_data_callback_id_;
 
  std::unordered_map<DataFrameCallbackId, RegisteredDataCallback>
      data_callbacks_;
 
  std::unordered_map<DataFrameCallbackId, std::shared_ptr<ActiveDataReader>>
      active_data_readers_;
 
  std::unordered_map<DataCallbackKey, std::shared_ptr<RemoteDataTrack>,
                     DataCallbackKeyHash>
      remote_data_tracks_;
 
  static constexpr int kMaxActiveReaders = 20;
};
 
} // namespace livekit
 
#endif /* LIVEKIT_SUBSCRIPTION_THREAD_DISPATCHER_H */