collector.h

//
// Copyright 2018 Pixar
//
// Licensed under the Apache License, Version 2.0 (the "Apache License")
// with the following modification; you may not use this file except in
// compliance with the Apache License and the following modification to it:
// Section 6. Trademarks. is deleted and replaced with:
//
// 6. Trademarks. This License does not grant permission to use the trade
//    names, trademarks, service marks, or product names of the Licensor
//    and its affiliates, except as required to comply with Section 4(c) of
//    the License and to reproduce the content of the NOTICE file.
//
// You may obtain a copy of the Apache License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the Apache License with the above modification is
// distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the Apache License for the specific
// language governing permissions and limitations under the Apache License.
//

#ifndef PXR_BASE_TRACE_COLLECTOR_H
#define PXR_BASE_TRACE_COLLECTOR_H

#include "pxr/pxr.h"

#include "pxr/base/trace/api.h"
#include "pxr/base/trace/concurrentList.h"
#include "pxr/base/trace/collection.h"
#include "pxr/base/trace/event.h"
#include "pxr/base/trace/key.h"
#include "pxr/base/trace/threads.h"

#include "pxr/base/tf/declarePtrs.h"
#include "pxr/base/tf/mallocTag.h"

#include "pxr/base/tf/pyTracing.h"

#include "pxr/base/tf/singleton.h"
#include "pxr/base/tf/refBase.h"
#include "pxr/base/tf/refPtr.h"
#include "pxr/base/tf/weakBase.h"
#include "pxr/base/tf/weakPtr.h"

#include <atomic>
#include <string>
#include <vector>

#include <tbb/spin_mutex.h>

PXR_NAMESPACE_OPEN_SCOPE

class TraceScopeHolder;
 
TF_DECLARE_WEAK_PTRS(TraceCollector);
TF_DECLARE_WEAK_AND_REF_PTRS(TraceScope);


class TraceCollector : public TfWeakBase {
public:

    TF_MALLOC_TAG_NEW("Trace", "TraceCollector");

    using This = TraceCollector;
    using ThisPtr = TraceCollectorPtr;

    using TimeStamp = TraceEvent::TimeStamp;

    using Key = TraceDynamicKey;

    TRACE_API static TraceCollector& GetInstance() {
         return TfSingleton<TraceCollector>::GetInstance();
    }
    
    TRACE_API ~TraceCollector();

    TRACE_API void SetEnabled(bool isEnabled);

    static bool IsEnabled() {
        return (_isEnabled.load(std::memory_order_acquire) == 1);
    }

    struct DefaultCategory {
        static constexpr TraceCategoryId GetId() { return TraceCategory::Default;}
        static bool IsEnabled() { return TraceCollector::IsEnabled(); }
    };

#ifdef PXR_PYTHON_SUPPORT_ENABLED
    bool IsPythonTracingEnabled() const {
        return _isPythonTracingEnabled.load(std::memory_order_acquire) != 0;
    }

    TRACE_API void SetPythonTracingEnabled(bool enabled);
#endif // PXR_PYTHON_SUPPORT_ENABLED

    TRACE_API TimeStamp GetScopeOverhead() const;
    
    TRACE_API void Clear();


    template <typename Category = DefaultCategory>
    TimeStamp BeginEvent(const Key& key) {
        if (ARCH_LIKELY(!Category::IsEnabled())) {
            return 0;
        } 
        return _BeginEvent(key, Category::GetId());
    }

    template <typename Category = DefaultCategory>
    void BeginEventAtTime(const Key& key, double ms) {
      if (ARCH_LIKELY(!Category::IsEnabled())) {
          return;
      }
      _BeginEventAtTime(key, ms, Category::GetId());
    }

    template <typename Category = DefaultCategory>
    TimeStamp EndEvent(const Key& key) {
        if (ARCH_LIKELY(!Category::IsEnabled())) {
            return 0;
        }
        return _EndEvent(key, Category::GetId());
    }

    template <typename Category = DefaultCategory>
    void EndEventAtTime(const Key& key, double ms) {
        if (ARCH_LIKELY(!Category::IsEnabled())) {
            return;
        }
        _EndEventAtTime(key, ms, Category::GetId());
    }



    template <typename Category = DefaultCategory>
    TimeStamp MarkerEvent(const Key& key) {
        if (ARCH_LIKELY(!Category::IsEnabled())) {
            return 0;
        } 
        return _MarkerEvent(key, Category::GetId());
    }

    template <typename Category = DefaultCategory>
    void MarkerEventAtTime(const Key& key, double ms) {
      if (ARCH_LIKELY(!Category::IsEnabled())) {
          return;
      }
      _MarkerEventAtTime(key, ms, Category::GetId());
    }

    template <typename Category = DefaultCategory>
    void BeginScope(const TraceKey& _key) {
        if (ARCH_LIKELY(!Category::IsEnabled()))
            return;

        _BeginScope(_key, Category::GetId());
    }

    template <typename Category, typename... Args>
    void BeginScope(
        const TraceKey& key, Args&&... args) {
        static_assert( sizeof...(Args) %2 == 0, 
            "Data arguments must come in pairs");

        if (ARCH_LIKELY(!Category::IsEnabled()))
            return;

        _PerThreadData *threadData = _GetThreadData();
        threadData->BeginScope(key, Category::GetId());
        _StoreDataRec(threadData, Category::GetId(), std::forward<Args>(args)...);
    }

    template <typename... Args>
    void BeginScope(const TraceKey& key, Args&&... args) {
        static_assert( sizeof...(Args) %2 == 0, 
            "Data arguments must come in pairs");

        // Explicitly cast to TraceCategoryId so overload resolution choose the
        // version with a category arguement.
        BeginScope<DefaultCategory>(key, 
            std::forward<Args>(args)...);
    }

    template <typename Category = DefaultCategory>
    void EndScope(const TraceKey& key) {
        if (ARCH_LIKELY(!Category::IsEnabled()))
            return;

        _EndScope(key, Category::GetId());
    }

    TRACE_API
    static void
    Scope(const TraceKey& key, TimeStamp start, TimeStamp stop) noexcept;

    template <typename Category = DefaultCategory>
    void Scope(const TraceKey& key, TimeStamp start, TimeStamp stop) {
        if (ARCH_LIKELY(!Category::IsEnabled()))
            return;
        _PerThreadData *threadData = _GetThreadData();
        threadData->EmplaceEvent(
            TraceEvent::Timespan, key, start, stop, Category::GetId());
    }
    
    template <typename Category, typename... Args>
    void ScopeArgs(Args&&... args) {
        static_assert( sizeof...(Args) %2 == 0, 
            "Data arguments must come in pairs");

        if (ARCH_LIKELY(!Category::IsEnabled()))
            return;

        _PerThreadData *threadData = _GetThreadData();
        _StoreDataRec(threadData, Category::GetId(), std::forward<Args>(args)...);
    }

    template <typename... Args>
    void ScopeArgs(Args&&... args) {
        static_assert( sizeof...(Args) %2 == 0, 
            "Data arguments must come in pairs");

        ScopeArgs<DefaultCategory>(std::forward<Args>(args)...);
    }


    template <typename Category = DefaultCategory>
    void MarkerEventStatic(const TraceKey& key) {
        if (ARCH_LIKELY(!Category::IsEnabled()))
            return;

        _PerThreadData *threadData = _GetThreadData();
        threadData->EmplaceEvent(
            TraceEvent::Marker, key, Category::GetId());
    }

    template <typename Category = DefaultCategory, typename T>
    void StoreData(const TraceKey &key, const T& value) {
        if (ARCH_UNLIKELY(Category::IsEnabled())) {
            _StoreData(_GetThreadData(), key, Category::GetId(), value);
        }
    }

    template <typename Category = DefaultCategory>
    void RecordCounterDelta(const TraceKey &key, 
                            double delta) {
        // Only record counter values if the collector is enabled.
        if (ARCH_UNLIKELY(Category::IsEnabled())) {
            _PerThreadData *threadData = _GetThreadData();
            threadData->EmplaceEvent(
                TraceEvent::CounterDelta, key, delta, Category::GetId());
        }
    }

    template <typename Category = DefaultCategory>
    void RecordCounterDelta(const Key &key, double delta) {
        if (ARCH_UNLIKELY(Category::IsEnabled())) {
            _PerThreadData *threadData = _GetThreadData();
            threadData->CounterDelta(key, delta, Category::GetId());
        }
    }

    template <typename Category = DefaultCategory>
    void RecordCounterValue(const TraceKey &key, double value) {
        // Only record counter values if the collector is enabled.
        if (ARCH_UNLIKELY(Category::IsEnabled())) {
            _PerThreadData *threadData = _GetThreadData();
            threadData->EmplaceEvent(
                TraceEvent::CounterValue, key, value, Category::GetId());
        }
    }

    template <typename Category = DefaultCategory>
    void RecordCounterValue(const Key &key, double value) {

        if (ARCH_UNLIKELY(Category::IsEnabled())) {
            _PerThreadData *threadData = _GetThreadData();
            threadData->CounterValue(key, value, Category::GetId());
        }
    }


    const std::string& GetLabel() {
        return _label;
    }

    TRACE_API void CreateCollection();

private:

    TraceCollector();

    friend class TfSingleton<TraceCollector>;

    class _PerThreadData;

    // Return a pointer to existing per-thread data or create one if none
    // exists.
    TRACE_API _PerThreadData* _GetThreadData() noexcept;

    TRACE_API TimeStamp _BeginEvent(const Key& key, TraceCategoryId cat);

    TRACE_API void _BeginEventAtTime(
        const Key& key, double ms, TraceCategoryId cat);

    TRACE_API TimeStamp _EndEvent(const Key& key, TraceCategoryId cat);

    TRACE_API void _EndEventAtTime(
         const Key& key, double ms, TraceCategoryId cat);

    TRACE_API TimeStamp _MarkerEvent(const Key& key, TraceCategoryId cat);

    TRACE_API void _MarkerEventAtTime(
        const Key& key, double ms, TraceCategoryId cat);

    // This is the fast execution path called from the TRACE_FUNCTION
    // and TRACE_SCOPE macros
    void _BeginScope(const TraceKey& key, TraceCategoryId cat)
    {
        // Note we're not calling _NewEvent, don't need to cache key
        _PerThreadData *threadData = _GetThreadData();
        threadData->BeginScope(key, cat);
    }

    // This is the fast execution path called from the TRACE_FUNCTION
    // and TRACE_SCOPE macros    
    TRACE_API void _EndScope(const TraceKey& key, TraceCategoryId cat);

    TRACE_API void _MeasureScopeOverhead();

#ifdef PXR_PYTHON_SUPPORT_ENABLED
    // Callback function registered as a python tracing function.
    void _PyTracingCallback(const TfPyTraceInfo &info);
#endif // PXR_PYTHON_SUPPORT_ENABLED

    // Implementation for small data that can stored inlined with the event.
    template <typename T,
        typename std::enable_if<
            sizeof(T) <= sizeof(uintptr_t) 
            && !std::is_pointer<T>::value , int>::type = 0>
    void _StoreData(_PerThreadData* threadData, const TraceKey &key,
        TraceCategoryId cat, const T& value) {
        threadData->StoreData(key, value, cat);
    }

    // Implementation for data that must be stored outside of the events.
    template <typename T,
        typename std::enable_if<
            (sizeof(T) > sizeof(uintptr_t))
            && !std::is_pointer<T>::value, int>::type = 0>
    void _StoreData(_PerThreadData* threadData, const TraceKey &key,
        TraceCategoryId cat,  const T& value) {
        threadData->StoreLargeData(key, value, cat);
    }

    // Specialization for c string
    void _StoreData(
        _PerThreadData* threadData, 
        const TraceKey &key, 
        TraceCategoryId cat, 
        const char* value) {
        threadData->StoreLargeData(key, value, cat);
    }

    // Specialization for std::string
    void _StoreData(
        _PerThreadData* threadData,
        const TraceKey &key,
        TraceCategoryId cat, 
        const std::string& value) {
        threadData->StoreLargeData(key, value.c_str(), cat);
    }

    // Variadic version to store multiple data events in one function call.
    template <typename K, typename T, typename... Args>
    void _StoreDataRec(
        _PerThreadData* threadData, TraceCategoryId cat, K&& key, 
        const T& value, Args&&... args) {
        _StoreData(threadData, std::forward<K>(key), cat, value);
        _StoreDataRec(threadData, cat, std::forward<Args>(args)...);
    }

    // Base case to terminate template recursion
    void _StoreDataRec(_PerThreadData* threadData, TraceCategoryId cat) {}


    // Thread-local storage, accessed via _GetThreadData()
    //
    class _PerThreadData {
        public:
            using EventList = TraceCollection::EventList;

            _PerThreadData();
            ~_PerThreadData();

            const TraceThreadId& GetThreadId() const {
                return _threadIndex;
            }
            TimeStamp BeginEvent(const Key& key, TraceCategoryId cat);
            TimeStamp EndEvent(const Key& key, TraceCategoryId cat);
            TimeStamp MarkerEvent(const Key& key, TraceCategoryId cat);

            // Debug Methods
            void BeginEventAtTime(
                const Key& key, double ms, TraceCategoryId cat);
            void EndEventAtTime(const Key& key, double ms, TraceCategoryId cat);
            void MarkerEventAtTime(const Key& key, double ms, TraceCategoryId cat);

            void BeginScope(const TraceKey& key, TraceCategoryId cat) {
                AtomicRef lock(_writing);
                _BeginScope(key, cat);
            }

            void EndScope(const TraceKey& key, TraceCategoryId cat) {
                AtomicRef lock(_writing);
                _EndScope(key, cat);
            }

            TRACE_API void CounterDelta(
                const Key&, double value, TraceCategoryId cat);

            TRACE_API void CounterValue(
                const Key&, double value, TraceCategoryId cat);

            template <typename T>
            void StoreData(
                const TraceKey& key, const T& data, TraceCategoryId cat) {
                AtomicRef lock(_writing);
                _events.load(std::memory_order_acquire)->EmplaceBack(
                    TraceEvent::Data, key, data, cat);
            }

            template <typename T>
            void StoreLargeData(
                const TraceKey& key, const T& data, TraceCategoryId cat) {
                AtomicRef lock(_writing);
                EventList* events = _events.load(std::memory_order_acquire);
                const auto* cached = events->StoreData(data);
                events->EmplaceBack(TraceEvent::Data, key, cached, cat);
            }

            template <typename... Args>
            void EmplaceEvent(Args&&... args) {
                AtomicRef lock(_writing);
                _events.load(std::memory_order_acquire)->EmplaceBack(
                    std::forward<Args>(args)...);
            }

#ifdef PXR_PYTHON_SUPPORT_ENABLED
            void PushPyScope(const Key& key, bool enabled);
            void PopPyScope(bool enabled);
#endif // PXR_PYTHON_SUPPORT_ENABLED

            // These methods can be called from threads at the same time as the 
            // other methods.
            std::unique_ptr<EventList> GetCollectionData();
            void Clear();

        private:
            void _BeginScope(const TraceKey& key, TraceCategoryId cat) {
                _events.load(std::memory_order_acquire)->EmplaceBack(
                    TraceEvent::Begin, key, cat);
            }

            void _EndScope(const TraceKey& key, TraceCategoryId cat);

            // Flag to let other threads know that the list is being written to.
            mutable std::atomic<bool> _writing;
            std::atomic<EventList*> _events;

            class AtomicRef {
            public:
                AtomicRef(std::atomic<bool>& b) : _bool(b) {
                    _bool.store(true, std::memory_order_release);
                }
                ~AtomicRef() {
                    _bool.store(false, std::memory_order_release);
                }
            private:
                std::atomic<bool>& _bool;
            };

            // An integer that is unique for each thread launched by any
            // threadDispatcher.  Each time a thread is Start-ed it get's
            // a new id.
            //
            TraceThreadId _threadIndex;

            // When auto-tracing python frames, this stores the stack of scopes.
            struct PyScope {
                Key key;
            };
            std::vector<PyScope> _pyScopes;
    };

    TRACE_API static std::atomic<int> _isEnabled;

    // A list with one _PerThreadData per thread.
    TraceConcurrentList<_PerThreadData> _allPerThreadData;

    std::string _label;

    TimeStamp _measuredScopeOverhead;

    std::atomic<int> _isPythonTracingEnabled;
    TfPyTraceFnId _pyTraceFnId;
};
 
TRACE_API_TEMPLATE_CLASS(TfSingleton<TraceCollector>);

PXR_NAMESPACE_CLOSE_SCOPE

#endif // PXR_BASE_TRACE_COLLECTOR_H