Commit 833a7af3 authored by Bob Van Landuyt's avatar Bob Van Landuyt 💬
Browse files

Add context with metadata and correlation 

The context can be passed from the client applications as a hash.

Every time a new context is added, it will inherit the values of the
previous context. Those previous values will be overridden by any new
ones passed.

A context always has a correlation id that is not empty. If an empty
correlation id is passed, a new one will be generated.

There is a `Labkit::Middleware::Rack` that will set the correlation id
to the current request id. Because of this, every newly started
context within a rack request will get the same correlation id.

Using the `Labkit::Middleware::Sidekiq::Client` middleware the context
active at the moment of scheduling a job will be serialized into the
job.

The `Labkit::Middleware::Sidekiq::Server` middleware will load the
context that was stored on the job, so newly scheduled jobs would
start with the same context. This context could be modified by the
job.

In this iteration the following extra values can be specified on the
context:
- user
- project
- root_namespace

**Usage of the Labkit::Context**

The preferred way of specifying a context is using a block:

  Labkit::Context.with_context(
    user: 'jane.doe',
    project: 'jane.doe/cool'
  ) do |context|
    # application code
  end

If there's no way to wrap the application code into a block (for
example for the grape api), then context can be pushed and popped:

  def context
    @context ||= Labkit::Context.push(user: 'jane.doe')
  end

  before { context }
  after { Labkit::Context.pop(context) }

It is possible to provide procs when assigning values to the context:

  Labkit::Context.with_context(user: -> { current_user }) { # ... }

This proc will only be executed when needed: When serialising the
context into a log or a job. This happens when calling
`Labkit::Context.to_h`.
parent 15809235

gitlab-labkit.gemspec

+2 −0
Original line number Diff line number Diff line
@@ -28,6 +28,7 @@ Gem::Specification.new do |spec| 


  # Please maintain alphabetical order for dev dependencies

  spec.add_development_dependency "grpc-tools", "~> 1.19"

  spec.add_development_dependency "pry", "~> 0.12"

  spec.add_development_dependency "rack", "~> 2.0"

  spec.add_development_dependency "rake", "~> 12.3"

  spec.add_development_dependency "rspec", "~> 3.8.0"
@@ -35,4 +36,5 @@ Gem::Specification.new do |spec| 
  spec.add_development_dependency "rubocop", "~> 0.65.0"

  spec.add_development_dependency "rubocop-rspec", "~> 1.22.1"

  spec.add_development_dependency "rufo", "~> 0.6"

  spec.add_development_dependency "sidekiq", "~> 5.2.7"

end

lib/gitlab-labkit.rb

+2 −0
Original line number Diff line number Diff line
@@ -8,8 +8,10 @@ require "active_support/all" 
# observability.

module Labkit

  autoload :Correlation, "labkit/correlation"

  autoload :Context, "labkit/context"

  autoload :Tracing, "labkit/tracing"

  autoload :Logging, "labkit/logging"

  autoload :Middleware, "labkit/middleware"

end



# rubocop:enable Naming/FileName

lib/labkit/context.rb

0 → 100644
+133 −0
Original line number Diff line number Diff line 
# frozen_string_literal: true



module Labkit

  # A context can be used to provide structured information on what resources

  # GitLab is working on within a service.

  # The currently supported keys are defined in the `KNOWN_KEYS` constant.

  #

  # Values can be provided by passing a hash. If one of the values is a Proc

  # the proc will only be called when the value is actually needed.

  #

  # Multiple contexts can be nested, the nested context will inherit the values

  # from the closest outer one.

  # All contexts will have the same correlation id.

  #

  # Usage:

  #   Labkit::Context.with_context(user: 'username', root_namespace: -> { get_root_namespace } do |context|

  #     logger.info(context.to_h)

  #   end

  #

  class Context

    LOG_KEY = "meta"

    CORRELATION_ID_KEY = "correlation_id"

    RAW_KEYS = [CORRELATION_ID_KEY].freeze

    KNOWN_KEYS = %w[user project root_namespace].freeze



    class << self

      def with_context(attributes = {})

        context = push(attributes)



        begin

          yield(context)

        ensure

          pop(context)

        end

      end



      def push(new_attributes = {})

        new_context = current&.merge(new_attributes) || new(new_attributes)



        contexts.push(new_context)



        new_context

      end



      def pop(context)

        contexts.pop while contexts.include?(context)

      end



      def correlation_id

        contexts.last&.correlation_id

      end



      def current

        contexts.last

      end



      def log_key(key)

        key = key.to_s

        return key if RAW_KEYS.include?(key)

        return key if key.starts_with?("#{LOG_KEY}.")



        "#{LOG_KEY}.#{key}"

      end



      def known_log_keys

        @known_log_keys ||= (KNOWN_KEYS.map(&method(:log_key)) + RAW_KEYS).freeze

      end



      private



      def contexts

        Thread.current[:labkit_contexts] ||= []

      end

    end



    def initialize(values = {})

      @data = {}



      assign_attributes(values)

    end



    def merge(new_attributes)

      new_context = self.class.new(data.dup)

      new_context.assign_attributes(new_attributes)



      new_context

    end



    def to_h

      expand_data

    end



    def correlation_id

      data[CORRELATION_ID_KEY]

    end



    protected



    def assign_attributes(attributes)

      attributes = attributes.transform_keys(&method(:log_key))

      attributes = attributes.slice(*known_log_keys)



      data.merge!(attributes)



      # Remove keys that had their values set to `nil` in the new attributes

      data.keep_if { |_, value| value.present? }



      # Assign a correlation if it was missing in the first context or when

      # explicitly removed

      data[CORRELATION_ID_KEY] ||= new_id



      data

    end



    private



    delegate :log_key, :known_log_keys, to: :class



    attr_reader :data



    def expand_data

      data.transform_values do |value|

        value = value.respond_to?(:call) ? value.call : value



        value.presence

      end.compact

    end



    def new_id

      SecureRandom.hex

    end

  end

end

lib/labkit/correlation/correlation_id.rb

+6 −23
Original line number Diff line number Diff line
@@ -5,38 +5,21 @@ module Labkit 
    # CorrelationId module provides access the Correlation-ID

    # of the current request

    module CorrelationId

      LOG_KEY = "correlation_id"

      LOG_KEY = Labkit::Context::CORRELATION_ID_KEY



      class << self

        def use_id(correlation_id, &_blk)

          # always generate a id if null is passed

          correlation_id ||= new_id



          ids.push(correlation_id || new_id)



          begin

            yield(current_id)

          ensure

            ids.pop

        def use_id(correlation_id)

          Labkit::Context.with_context(LOG_KEY => correlation_id) do |context|

            yield(context.correlation_id)

          end

        end



        def current_id

          ids.last

          Labkit::Context.correlation_id

        end



        def current_or_new_id

          current_id || new_id

        end



        private



        def ids

          Thread.current[:correlation_id] ||= []

        end



        def new_id

          SecureRandom.uuid

          current_id || Labkit::Context.push.correlation_id

        end

      end

    end

lib/labkit/middleware.rb

0 → 100644
+9 −0
Original line number Diff line number Diff line 
# frozen_string_literal: true



module Labkit

  # Adds middlewares for using in rack and sidekiq

  module Middleware

    autoload :Rack, "labkit/middleware/rack"

    autoload :Sidekiq, "labkit/middleware/sidekiq"

  end

end