frostfs-observability/logging/lokicore/core.go

package lokicore

import (
	"time"

	"git.frostfs.info/TrueCloudLab/frostfs-observability/logging/lokicore/loki"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
)

// Zap Core for loki.
// Expands the zapcore.Core interface with calls to export logs to Loki.
type LokiCore struct {
	original zapcore.Core
	encoder  zapcore.Encoder
	loki     *loki.Client
}

func New(original zapcore.Core, lokiCfg loki.Config) *LokiCore {
	encoderConfig := zap.NewProductionEncoderConfig()
	encoder := zapcore.NewJSONEncoder(encoderConfig)
	encoderConfig.EncodeTime = zapcore.TimeEncoderOfLayout(time.RFC3339Nano)

	return &LokiCore{
		original: original,
		encoder:  encoder,
		loki:     loki.Setup(lokiCfg),
	}
}

func (c *LokiCore) With(fields []zapcore.Field) zapcore.Core {
	return &LokiCore{
		original: c.original.With(fields),
		encoder:  c.encoder,
		loki:     c.loki,
	}
}

func (c *LokiCore) Write(entry zapcore.Entry, fields []zapcore.Field) error {
	if err := c.original.Write(entry, fields); err != nil {
		return err
	}

	buffer, err := c.encoder.EncodeEntry(entry, fields)
	defer buffer.Free()

	if err != nil {
		return err
	}

	return c.loki.Send(buffer.String(), entry.Time)
}

func (c *LokiCore) Check(entry zapcore.Entry, checked *zapcore.CheckedEntry) *zapcore.CheckedEntry {
	if c.Enabled(entry.Level) {
		return checked.AddCore(entry, c)
	}
	return checked
}

func (c *LokiCore) Sync() error {
	return c.original.Sync()
}

func (c *LokiCore) Enabled(level zapcore.Level) bool {
	return c.original.Enabled(level)
}