path: root/library/cpp/lwtrace/control.h

#pragma once 
 
#include "custom_action.h"
#include "event.h"
#include "log.h"
#include "log_shuttle.h"
#include "probe.h"

#include <library/cpp/lwtrace/protos/lwtrace.pb.h>

#include <google/protobuf/repeated_field.h>

#include <util/generic/deque.h> 
#include <util/generic/hash.h> 
#include <util/generic/noncopyable.h> 
#include <util/generic/ptr.h> 
#include <util/generic/set.h> 
#include <util/generic/vector.h> 
 
namespace NLWTrace { 

    using TProbeMap = THashMap<std::pair<TString, TString>, TProbe*>;

    // Interface for probe ownership management 
    class IBox: public virtual TThrRefBase { 
    private: 
        bool Owns; 

    public: 
        explicit IBox(bool ownsProbe = false) 
            : Owns(ownsProbe) 
        {
        }
 
        bool OwnsProbe() { 
            return Owns; 
        } 
 
        virtual TProbe* GetProbe() = 0; 
    }; 
 
    using TBoxPtr = TIntrusivePtr<IBox>; 
 
    // Simple not-owning box, that just holds pointer to static/global probe (e.g. created by LWTRACE_DEFINE_PROVIDER) 
    class TStaticBox: public IBox { 
    private: 
        TProbe* Probe; 

    public: 
        explicit TStaticBox(TProbe* probe) 
            : IBox(false) 
            , Probe(probe) 
        {
        }
 
        TProbe* GetProbe() override { 
            return Probe; 
        } 
    }; 
 
    // Just a set of unique probes 
    // TODO[serxa]: get rid of different ProbeRegistries, use unique one (singleton) with auto registration 
    class TProbeRegistry: public TNonCopyable { 
    private: 
        TMutex Mutex; 
 
        // Probe* pointer uniquely identifies a probe and boxptr actually owns TProbe object (if required) 
        using TProbes = THashMap<TProbe*, TBoxPtr>; 
        TProbes Probes; 
 
        // Probe provider-name pairs must be unique, keep track of them 
        using TIds = TSet<std::pair<TString, TString>>; 
        TIds Ids; 
 
    public: 
        // Add probes from null-terminated array of probe pointers. Probe can be added multiple times. Thread-safe. 
        // Implies probes you pass will live forever (e.g. created by LWTRACE_DEFINE_PROVIDER) 
        void AddProbesList(TProbe** reg); 
 
        // Manage probes that are created/destructed dynamically 
        void AddProbe(const TBoxPtr& box); 
        void RemoveProbe(TProbe* probe); 
 
        // Helper class to make thread-safe iteration over probes 
        class TProbesAccessor { 
        private: 
            TGuard<TMutex> Guard; 
            TProbes& Probes; 
        public: 
            explicit TProbesAccessor(TProbeRegistry* registry) 
                : Guard(registry->Mutex) 
                , Probes(registry->Probes) 
            {} 
 
            explicit TProbesAccessor(TProbeRegistry& registry) 
                : TProbesAccessor(&registry) 
            {} 
 
            auto begin() { return Probes.begin(); } 
            auto end() { return Probes.end(); } 
        }; 
 
        friend class TProbesAccessor; 
 
    private: 
        void AddProbeNoLock(const TBoxPtr& box); 
        void RemoveProbeNoLock(TProbe* probe); 
    }; 
 
    // Represents a compiled trace query, holds executors attached to probes 
    class TSession: public TNonCopyable { 
    public: 
        typedef THashMap<TString, TAtomicBase> TTraceVariables; 
 
    private: 
        const TInstant StartTime; 
        const ui64 TraceIdx; 
        TProbeRegistry& Registry; 
        TDuration StoreDuration; 
        TDuration ReadDuration; 
        TCyclicLog CyclicLog; 
        TDurationLog DurationLog; 
        TCyclicDepot CyclicDepot; 
        TDurationDepot DurationDepot; 
        TAtomic LastTrackId; 
        TAtomic LastSpanId;
        typedef TVector<std::pair<TProbe*, IExecutor*>> TProbes; 
        TProbes Probes; 
        bool Attached; 
        NLWTrace::TQuery Query; 
        TTraceVariables TraceVariables; 
        TTraceResources TraceResources; 
        void InsertExecutor(TTraceVariables& traceVariables, 
                            size_t bi, const NLWTrace::TPredicate* pred, 
                            const google::protobuf::RepeatedPtrField<NLWTrace::TAction>& actions, 
                            TProbe* probe, const bool destructiveActionsAllowed, 
                            const TCustomActionFactory& customActionFactory); 
 
    private: 
        void Destroy(); 
 
    public: 
        TSession(ui64 traceIdx, 
               TProbeRegistry& registry, 
               const NLWTrace::TQuery& query, 
               const bool destructiveActionsAllowed, 
               const TCustomActionFactory& customActionFactory); 
        ~TSession(); 
        void Detach(); 
        size_t GetEventsCount() const; 
        size_t GetThreadsCount() const; 
        const NLWTrace::TQuery& GetQuery() const { 
            return Query; 
        } 
        TInstant GetStartTime() const { 
            return StartTime; 
        } 
        ui64 GetTraceIdx() const { 
            return TraceIdx; 
        } 
        TTraceResources& Resources() { 
            return TraceResources; 
        } 
        const TTraceResources& Resources() const { 
            return TraceResources; 
        } 
 
        template <class TReader> 
        void ReadThreads(TReader& r) const { 
            CyclicLog.ReadThreads(r); 
            DurationLog.ReadThreads(r); 
        } 
 
        template <class TReader> 
        void ReadItems(TReader& r) const { 
            CyclicLog.ReadItems(r); 
            DurationLog.ReadItems(GetCycleCount(), DurationToCycles(ReadDuration), r); 
        } 
 
        template <class TReader> 
        void ReadItems(ui64 now, ui64 duration, TReader& r) const { 
            CyclicLog.ReadItems(r); 
            DurationLog.ReadItems(now, duration, r); 
        } 
 
        template <class TReader> 
        void ReadDepotThreads(TReader& r) const { 
            CyclicDepot.ReadThreads(r); 
            DurationDepot.ReadThreads(r); 
        } 
 
        template <class TReader> 
        void ReadDepotItems(TReader& r) const { 
            CyclicDepot.ReadItems(r); 
            DurationDepot.ReadItems(GetCycleCount(), DurationToCycles(ReadDuration), r); 
        } 
 
        template <class TReader> 
        void ReadDepotItems(ui64 now, ui64 duration, TReader& r) const { 
            CyclicDepot.ReadItems(r); 
            DurationDepot.ReadItems(now, duration, r); 
        } 
        void ToProtobuf(TLogPb& pb) const; 
    }; 
 
    // Deserialization result.
    // Either IsSuccess is true or FailedEventNames contains event names
    // we were not able to deserialize.
    struct TTraceDeserializeStatus
    {
        bool IsSuccess = true;
        TVector<TString> FailedEventNames;

        void AddFailedEventName(const TString& name)
        {
            IsSuccess = false;
            FailedEventNames.emplace_back(name);
        }
    };

    // Just a registry of all active trace queries 
    // Facade for all interactions with probes/traces 
    class TManager: public TNonCopyable { 
    private: 
        TProbeRegistry& Registry; 
        TMutex Mtx; 
        ui64 LastTraceIdx = 1;
        typedef THashMap<TString, TSession*> TTraces; // traceId -> TSession 
        TTraces Traces; 
        bool DestructiveActionsAllowed; 
        TCustomActionFactory CustomActionFactory; 
        THolder<TRunLogShuttleActionExecutor<TCyclicDepot>> SerializingExecutor;
 
    public: 
        static constexpr ui64 RemoteTraceIdx = 0;

    public:
        TManager(TProbeRegistry& registry, bool allowDestructiveActions); 
        ~TManager(); 
        bool HasTrace(const TString& id) const; 
        const TSession* GetTrace(const TString& id) const; 
        void New(const TString& id, const NLWTrace::TQuery& query); 
        void Delete(const TString& id); 
        void Stop(const TString& id); 
 
        template <class TReader> 
        void ReadProbes(TReader& reader) const { 
            TProbeRegistry::TProbesAccessor probes(Registry); 
            for (auto& kv : probes) { 
                TProbe* probe = kv.first; 
                reader.Push(probe); 
            } 
        } 

        template <class TReader> 
        void ReadTraces(TReader& reader) const { 
            TGuard<TMutex> g(Mtx); 
            for (const auto& Trace : Traces) { 
                const TString& id = Trace.first; 
                const TSession* trace = Trace.second; 
                reader.Push(id, trace); 
            } 
        } 
 
        template <class TReader> 
        void ReadLog(const TString& id, TReader& reader) { 
            TGuard<TMutex> g(Mtx); 
            TTraces::iterator it = Traces.find(id); 
            if (it == Traces.end()) { 
                ythrow yexception() << "trace id '" << id << "' is not used"; 
            } else { 
                TSession* trace = it->second; 
                trace->ReadItems(reader); 
            } 
        } 
 
        template <class TReader> 
        void ReadLog(const TString& id, ui64 now, TReader& reader) { 
            TGuard<TMutex> g(Mtx); 
            TTraces::iterator it = Traces.find(id); 
            if (it == Traces.end()) { 
                ythrow yexception() << "trace id '" << id << "' is not used"; 
            } else { 
                TSession* trace = it->second; 
                trace->ReadItems(now, reader); 
            } 
        } 
 
        template <class TReader> 
        void ReadDepot(const TString& id, TReader& reader) { 
            TGuard<TMutex> g(Mtx); 
            TTraces::iterator it = Traces.find(id); 
            if (it == Traces.end()) { 
                ythrow yexception() << "trace id '" << id << "' is not used"; 
            } else { 
                TSession* trace = it->second; 
                trace->ReadDepotItems(reader); 
            } 
        } 
 
        template <class TReader> 
        void ReadDepot(const TString& id, ui64 now, TReader& reader) { 
            TGuard<TMutex> g(Mtx); 
            TTraces::iterator it = Traces.find(id); 
            if (it == Traces.end()) { 
                ythrow yexception() << "trace id '" << id << "' is not used"; 
            } else { 
                TSession* trace = it->second; 
                trace->ReadDepotItems(now, reader); 
            } 
        } 
 
        bool GetDestructiveActionsAllowed() { 
            return DestructiveActionsAllowed; 
        } 
 
        void RegisterCustomAction(const TString& name, const TCustomActionFactory::TCallback& callback) { 
            CustomActionFactory.Register(name, callback); 
        } 
 
        template <class T> 
        void RegisterCustomAction() { 
            CustomActionFactory.Register(T::GetActionName(), [=](TProbe* probe, const TCustomAction& action, TSession* trace) { 
                return new T(probe, action, trace); 
            }); 
        } 

        TProbeMap GetProbesMap();

        void CreateTraceRequest(TTraceRequest& msg, TOrbit& orbit);

        bool HandleTraceRequest(
            const TTraceRequest& msg,
            TOrbit& orbit);

        TTraceDeserializeStatus HandleTraceResponse(
            const TTraceResponse& msg,
            const TProbeMap& probesMap,
            TOrbit& orbit,
            i64 timeOffset = 0,
            double timeScale = 1);

        void CreateTraceResponse(
            TTraceResponse& msg,
            TOrbit& orbit);

        bool IsTraced(TOrbit& orbit) {
            return orbit.HasShuttle(TManager::RemoteTraceIdx);
        }
    }; 
}