go.opentelemetry.io/contrib/instrumentation/github.com/Shopify/sarama/otelsarama@v0.43.0/message.go (about)

     1  // Copyright The OpenTelemetry Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package otelsarama // import "go.opentelemetry.io/contrib/instrumentation/github.com/Shopify/sarama/otelsarama"
    16  
    17  import (
    18  	"github.com/Shopify/sarama"
    19  
    20  	"go.opentelemetry.io/otel/propagation"
    21  )
    22  
    23  var _ propagation.TextMapCarrier = (*ProducerMessageCarrier)(nil)
    24  var _ propagation.TextMapCarrier = (*ConsumerMessageCarrier)(nil)
    25  
    26  // ProducerMessageCarrier injects and extracts traces from a sarama.ProducerMessage.
    27  type ProducerMessageCarrier struct {
    28  	msg *sarama.ProducerMessage
    29  }
    30  
    31  // NewProducerMessageCarrier creates a new ProducerMessageCarrier.
    32  func NewProducerMessageCarrier(msg *sarama.ProducerMessage) ProducerMessageCarrier {
    33  	return ProducerMessageCarrier{msg: msg}
    34  }
    35  
    36  // Get retrieves a single value for a given key.
    37  func (c ProducerMessageCarrier) Get(key string) string {
    38  	for _, h := range c.msg.Headers {
    39  		if string(h.Key) == key {
    40  			return string(h.Value)
    41  		}
    42  	}
    43  	return ""
    44  }
    45  
    46  // Set sets a header.
    47  func (c ProducerMessageCarrier) Set(key, val string) {
    48  	// Ensure uniqueness of keys
    49  	for i := 0; i < len(c.msg.Headers); i++ {
    50  		if string(c.msg.Headers[i].Key) == key {
    51  			c.msg.Headers = append(c.msg.Headers[:i], c.msg.Headers[i+1:]...)
    52  			i--
    53  		}
    54  	}
    55  	c.msg.Headers = append(c.msg.Headers, sarama.RecordHeader{
    56  		Key:   []byte(key),
    57  		Value: []byte(val),
    58  	})
    59  }
    60  
    61  // Keys returns a slice of all key identifiers in the carrier.
    62  func (c ProducerMessageCarrier) Keys() []string {
    63  	out := make([]string, len(c.msg.Headers))
    64  	for i, h := range c.msg.Headers {
    65  		out[i] = string(h.Key)
    66  	}
    67  	return out
    68  }
    69  
    70  // ConsumerMessageCarrier injects and extracts traces from a sarama.ConsumerMessage.
    71  type ConsumerMessageCarrier struct {
    72  	msg *sarama.ConsumerMessage
    73  }
    74  
    75  // NewConsumerMessageCarrier creates a new ConsumerMessageCarrier.
    76  func NewConsumerMessageCarrier(msg *sarama.ConsumerMessage) ConsumerMessageCarrier {
    77  	return ConsumerMessageCarrier{msg: msg}
    78  }
    79  
    80  // Get retrieves a single value for a given key.
    81  func (c ConsumerMessageCarrier) Get(key string) string {
    82  	for _, h := range c.msg.Headers {
    83  		if h != nil && string(h.Key) == key {
    84  			return string(h.Value)
    85  		}
    86  	}
    87  	return ""
    88  }
    89  
    90  // Set sets a header.
    91  func (c ConsumerMessageCarrier) Set(key, val string) {
    92  	// Ensure uniqueness of keys
    93  	for i := 0; i < len(c.msg.Headers); i++ {
    94  		if c.msg.Headers[i] != nil && string(c.msg.Headers[i].Key) == key {
    95  			c.msg.Headers = append(c.msg.Headers[:i], c.msg.Headers[i+1:]...)
    96  			i--
    97  		}
    98  	}
    99  	c.msg.Headers = append(c.msg.Headers, &sarama.RecordHeader{
   100  		Key:   []byte(key),
   101  		Value: []byte(val),
   102  	})
   103  }
   104  
   105  // Keys returns a slice of all key identifiers in the carrier.
   106  func (c ConsumerMessageCarrier) Keys() []string {
   107  	out := make([]string, len(c.msg.Headers))
   108  	for i, h := range c.msg.Headers {
   109  		out[i] = string(h.Key)
   110  	}
   111  	return out
   112  }