git.zd.zone/hrpc/hrpc@v0.0.12/log/hook/kafka.go (about)

     1  package hook
     2  
     3  import (
     4  	"encoding/json"
     5  	"runtime"
     6  	"time"
     7  
     8  	"github.com/Shopify/sarama"
     9  	"github.com/sirupsen/logrus"
    10  )
    11  
    12  type kafka struct {
    13  	user, password string
    14  	endpoint       string
    15  	topic          string
    16  
    17  	producer sarama.SyncProducer
    18  }
    19  
    20  func NewKafka(user, password, endpoint, topic string) *kafka {
    21  	return &kafka{
    22  		user:     user,
    23  		password: password,
    24  		endpoint: endpoint,
    25  		topic:    topic,
    26  	}
    27  }
    28  
    29  func (k kafka) Fire(entry *logrus.Entry) error {
    30  	if _, _, err := k.producer.SendMessage(&sarama.ProducerMessage{
    31  		Topic: k.topic,
    32  		Value: sarama.StringEncoder(getContent(entry)),
    33  	}); err != nil {
    34  		return err
    35  	}
    36  	return nil
    37  }
    38  
    39  func (k kafka) Levels() []logrus.Level {
    40  	return []logrus.Level{
    41  		logrus.PanicLevel,
    42  		logrus.FatalLevel,
    43  		logrus.ErrorLevel,
    44  		logrus.WarnLevel,
    45  		logrus.InfoLevel,
    46  	}
    47  }
    48  
    49  func (k *kafka) Establish() error {
    50  	producer, err := sarama.NewSyncProducer([]string{k.endpoint}, k.config())
    51  	if err != nil {
    52  		return err
    53  	}
    54  	k.producer = producer
    55  	return nil
    56  }
    57  
    58  func (k kafka) config() *sarama.Config {
    59  	config := sarama.NewConfig()
    60  	config.Net.SASL.Mechanism = "PLAIN"
    61  	config.Net.SASL.Version = int16(1)
    62  	config.Net.SASL.Enable = true
    63  	config.Net.SASL.User = k.user
    64  	config.Net.SASL.Password = k.password
    65  	config.Producer.Return.Successes = true
    66  	config.Version = sarama.V0_11_0_0
    67  	return config
    68  }
    69  
    70  func getContent(entry *logrus.Entry) string {
    71  	var out = struct {
    72  		Fields  logrus.Fields
    73  		Level   string
    74  		Message string
    75  		Time    time.Time
    76  		Caller  *runtime.Frame
    77  	}{
    78  		Fields:  entry.Data,
    79  		Level:   entry.Level.String(),
    80  		Message: entry.Message,
    81  		Time:    entry.Time,
    82  		Caller:  entry.Caller,
    83  	}
    84  	d, err := json.Marshal(out)
    85  	if err != nil {
    86  		return ""
    87  	}
    88  	return string(d)
    89  }