content: /*
   * Copyright The OpenTelemetry Authors
   * SPDX-License-Identifier: Apache-2.0
   */
  
  package io.opentelemetry.instrumentation.rocketmqclient.v4_8;
  
  import io.opentelemetry.context.Context;
  import io.opentelemetry.instrumentation.api.util.VirtualField;
  import org.apache.rocketmq.client.hook.ConsumeMessageContext;
  import org.apache.rocketmq.client.hook.ConsumeMessageHook;
  
  final class TracingConsumeMessageHookImpl implements ConsumeMessageHook {
  
    private static final VirtualField<ConsumeMessageContext, ContextAndScope> contextAndScopeField =
        VirtualField.find(ConsumeMessageContext.class, ContextAndScope.class);
  
    private final RocketMqConsumerInstrumenter instrumenter;
  
    TracingConsumeMessageHookImpl(RocketMqConsumerInstrumenter instrumenter) {
      this.instrumenter = instrumenter;
    }
  
    @Override
    public String hookName() {
      return "OpenTelemetryConsumeMessageTraceHook";
    }
  
    @Override
    public void consumeMessageBefore(ConsumeMessageContext context) {
      if (context == null || context.getMsgList() == null || context.getMsgList().isEmpty()) {
        return;
      }
      Context parentContext = Context.current();
      Context newContext = instrumenter.start(parentContext, context.getMsgList());
  
      // it's safe to store the scope in the rocketMq message context, both before() and after()
      // methods are always called from the same thread; see:
      // - ConsumeMessageConcurrentlyService$ConsumeRequest#run()
      // - ConsumeMessageOrderlyService$ConsumeRequest#run()
      if (newContext != parentContext) {
        contextAndScopeField.set(
            context, ContextAndScope.create(newContext, newContext.makeCurrent()));
      }
    }
  
    @Override
    public void consumeMessageAfter(ConsumeMessageContext context) {
      if (context == null || context.getMsgList() == null || context.getMsgList().isEmpty()) {
        return;
      }
      ContextAndScope contextAndScope = contextAndScopeField.get(context);
      if (contextAndScope != null) {
        contextAndScope.close();
        instrumenter.end(contextAndScope.getContext(), context.getMsgList());
      }
    }
  }
  