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.
 */
 
#pragma once
 
#include <cstdint>
#include <functional>
#include <memory>
#include <mutex>
#include <optional>
#include <string>
#include <thread>
#include <unordered_map>
#include <vector>
 
#include "livekit/audio_stream.h"
#include "livekit/video_stream.h"
#include "livekit/visibility.h"
 
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 VideoFrameEventCallback = std::function<void(const VideoFrameEvent&)>;
 
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 LIVEKIT_API SubscriptionThreadDispatcher {
public:
  SubscriptionThreadDispatcher();
 
  ~SubscriptionThreadDispatcher();
 
  void setOnAudioFrameCallback(const std::string& participant_identity, const std::string& track_name,
                               AudioFrameCallback callback, const AudioStream::Options& opts = {});
 
  void setOnVideoFrameCallback(const std::string& participant_identity, const std::string& track_name,
                               VideoFrameCallback callback, const VideoStream::Options& opts = {});
 
  void setOnVideoFrameEventCallback(const std::string& participant_identity, const std::string& track_name,
                                    VideoFrameEventCallback callback, const VideoStream::Options& opts = {});
 
  void clearOnAudioFrameCallback(const std::string& participant_identity, const std::string& track_name);
 
  void clearOnVideoFrameCallback(const std::string& participant_identity, const std::string& track_name);
 
  void handleTrackSubscribed(const std::string& participant_identity, 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;
    std::string track_name;
 
    bool operator==(const CallbackKey& o) const {
      return participant_identity == o.participant_identity && 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<std::string>{}(k.track_name);
      return h1 ^ (h2 << 1);
    }
  };
 
  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 legacy_callback;
    VideoFrameEventCallback event_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,
                                     const AudioFrameCallback& cb, const AudioStream::Options& opts);
 
  std::thread startVideoReaderLocked(const CallbackKey& key, const std::shared_ptr<Track>& track,
                                     const RegisteredVideoCallback& callback);
 
  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, const 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_{0};
 
  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