processor.go

// Copyright The OpenTelemetry Authors
// SPDX-License-Identifier: Apache-2.0

package lsmintervalprocessor // import "github.com/open-telemetry/opentelemetry-collector-contrib/processor/lsmintervalprocessor"

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/cockroachdb/pebble"
	"github.com/cockroachdb/pebble/vfs"
	"go.opentelemetry.io/collector/component"
	"go.opentelemetry.io/collector/consumer"
	"go.opentelemetry.io/collector/pdata/pmetric"
	"go.opentelemetry.io/collector/processor"
	"go.uber.org/zap"

	"github.com/open-telemetry/opentelemetry-collector-contrib/processor/lsmintervalprocessor/internal/merger"
)

var _ processor.Metrics = (*Processor)(nil)
var zeroTime = time.Unix(0, 0).UTC()

// TODO: Optimize pebble
const batchCommitThreshold = 16 << 20 // 16MB

type Processor struct {
	db      *pebble.DB
	dataDir string
	dbOpts  *pebble.Options
	wOpts   *pebble.WriteOptions

	intervals []time.Duration
	next      consumer.Metrics

	mu             sync.Mutex
	batch          *pebble.Batch
	processingTime time.Time

	ctx           context.Context
	cancel        context.CancelFunc
	exportStopped chan struct{}
	logger        *zap.Logger
}

func newProcessor(cfg *Config, log *zap.Logger, next consumer.Metrics) (*Processor, error) {
	dbOpts := &pebble.Options{
		Merger: &pebble.Merger{
			Name: "pmetrics_merger",
			Merge: func(key, value []byte) (pebble.ValueMerger, error) {
				var v merger.Value
				if err := v.UnmarshalProto(value); err != nil {
					return nil, fmt.Errorf("failed to unmarshal value from db: %w", err)
				}
				return merger.New(v), nil
			},
		},
	}
	dataDir := cfg.Directory
	writeOpts := pebble.Sync
	if dataDir == "" {
		log.Info("no directory specified, switching to in-memory mode")
		dbOpts.FS = vfs.NewMem()
		dbOpts.DisableWAL = true
		writeOpts = pebble.NoSync
		dataDir = "/data" // will be created in the in-mem file-system
	}

	ctx, cancel := context.WithCancel(context.Background())
	return &Processor{
		dataDir:        dataDir,
		dbOpts:         dbOpts,
		wOpts:          writeOpts,
		intervals:      cfg.Intervals,
		next:           next,
		processingTime: time.Now().UTC().Truncate(cfg.Intervals[0]),
		ctx:            ctx,
		cancel:         cancel,
		logger:         log,
	}, nil
}

func (p *Processor) Start(ctx context.Context, host component.Host) error {
	p.mu.Lock()
	if p.db == nil {
		db, err := pebble.Open(p.dataDir, p.dbOpts)
		if err != nil {
			return fmt.Errorf("failed to open database: %w", err)
		}
		p.db = db
	}

	if p.exportStopped == nil {
		p.exportStopped = make(chan struct{})
	}
	p.mu.Unlock()

	go func() {
		defer close(p.exportStopped)
		to := p.processingTime.Add(p.intervals[0])
		timer := time.NewTimer(time.Until(to))
		defer timer.Stop()

		for {
			select {
			case <-p.ctx.Done():
				return
			case <-timer.C:
			}

			p.mu.Lock()
			batch := p.batch
			p.batch = nil
			p.processingTime = to
			p.mu.Unlock()

			// Export the batch
			if err := p.commitAndExport(p.ctx, batch, to); err != nil {
				p.logger.Warn("failed to export", zap.Error(err), zap.Time("end_time", to))
			}

			to = to.Add(p.intervals[0])
			timer.Reset(time.Until(to))
		}
	}()
	return nil
}

func (p *Processor) Shutdown(ctx context.Context) error {
	// Signal stop for the exporting goroutine
	p.cancel()

	// Wait for the exporting goroutine to stop. Note that we don't need to acquire
	// mutex here since even if there is a race between Start and Shutdown the
	// processor context is cancelled ensuring export goroutine will be noop.
	if p.exportStopped != nil {
		select {
		case <-ctx.Done():
			return fmt.Errorf("failed to shutdown due to context timeout while waiting for export to stop: %w", ctx.Err())
		case <-p.exportStopped:
		}
	}

	p.mu.Lock()
	defer p.mu.Unlock()

	// Ensure all data in the database is exported
	if p.db != nil {
		p.logger.Info("exporting all data before shutting down")
		if p.batch != nil {
			if err := p.batch.Commit(p.wOpts); err != nil {
				return fmt.Errorf("failed to commit batch: %w", err)
			}
			if err := p.batch.Close(); err != nil {
				return fmt.Errorf("failed to close batch: %w", err)
			}
			p.batch = nil
		}

		var errs []error
		for _, ivl := range p.intervals {
			// At any particular time there will be 1 export candidate for
			// each aggregation interval. We will align the end time and
			// process each of these.
			to := p.processingTime.Truncate(ivl).Add(ivl)
			if err := p.export(ctx, to); err != nil {
				errs = append(errs, fmt.Errorf(
					"failed to export metrics for interval %s: %w", ivl, err),
				)
			}
		}
		if len(errs) > 0 {
			return fmt.Errorf("failed while running final export: %w", errors.Join(errs...))
		}
		if err := p.db.Close(); err != nil {
			return fmt.Errorf("failed to close database: %w", err)
		}
		// All future operations are invalid after db is closed
		p.db = nil
	}
	return nil
}

func (p *Processor) Capabilities() consumer.Capabilities {
	return consumer.Capabilities{MutatesData: true}
}

func (p *Processor) ConsumeMetrics(ctx context.Context, md pmetric.Metrics) error {
	var errs []error
	v := merger.Value{Metrics: pmetric.NewMetrics()}
	md.ResourceMetrics().RemoveIf(func(rm pmetric.ResourceMetrics) bool {
		rm.ScopeMetrics().RemoveIf(func(sm pmetric.ScopeMetrics) bool {
			sm.Metrics().RemoveIf(func(m pmetric.Metric) bool {
				switch t := m.Type(); t {
				case pmetric.MetricTypeEmpty, pmetric.MetricTypeGauge, pmetric.MetricTypeSummary:
					return false
				case pmetric.MetricTypeSum, pmetric.MetricTypeHistogram:
					v.MergeMetric(rm, sm, m)
					return true
				case pmetric.MetricTypeExponentialHistogram:
					// TODO implement for parity with intervalprocessor
					return false
				default:
					// All metric types are handled, this is unexpected
					errs = append(errs, fmt.Errorf("unexpected metric type, dropping: %d", t))
					return true
				}
			})
			return sm.Metrics().Len() == 0
		})
		return rm.ScopeMetrics().Len() == 0
	})

	vb, err := v.MarshalProto()
	if err != nil {
		return errors.Join(append(errs, fmt.Errorf("failed to marshal value to proto binary: %w", err))...)
	}

	p.mu.Lock()
	defer p.mu.Unlock()

	keys := make([][]byte, len(p.intervals))
	for i, ivl := range p.intervals {
		key := merger.NewKey(ivl, p.processingTime)
		keys[i], err = key.Marshal()
		if err != nil {
			errs = append(errs, fmt.Errorf("failed to marshal key to binary for ivl %s: %w", ivl, err))
			continue
		}
	}

	if p.batch == nil {
		// TODO possible optimization by using NewBatchWithSize
		p.batch = p.db.NewBatch()
	}

	for _, k := range keys {
		if err := p.batch.Merge(k, vb, nil); err != nil {
			errs = append(errs, fmt.Errorf("failed to merge to db: %w", err))
		}
	}

	if p.batch.Len() >= batchCommitThreshold {
		if err := p.batch.Commit(p.wOpts); err != nil {
			return errors.Join(append(errs, fmt.Errorf("failed to commit a batch to db: %w", err))...)
		}
		if err := p.batch.Close(); err != nil {
			return errors.Join(append(errs, fmt.Errorf("failed to close a batch post commit: %w", err))...)
		}
		p.batch = nil
	}

	// Call next for the metrics remaining in the input
	if err := p.next.ConsumeMetrics(ctx, md); err != nil {
		errs = append(errs, err)
	}

	if len(errs) > 0 {
		return errors.Join(errs...)
	}
	return nil
}

// commitAndExport commits the batch to DB and exports all aggregated metrics in the provided range
// bounded by `to. If the batch is not committed then a corresponding error would be returned however
// exports will still proceed.
func (p *Processor) commitAndExport(ctx context.Context, batch *pebble.Batch, to time.Time) error {
	var errs []error
	if batch != nil {
		if err := batch.Commit(p.wOpts); err != nil {
			errs = append(errs, fmt.Errorf("failed to commit batch before export: %w", err))
		}
		if err := batch.Close(); err != nil {
			errs = append(errs, fmt.Errorf("failed to close batch before export: %w", err))
		}
	}
	if err := p.export(ctx, to); err != nil {
		errs = append(errs, fmt.Errorf("failed to export: %w", err))
	}
	if len(errs) > 0 {
		return errors.Join(errs...)
	}
	return nil
}

func (p *Processor) export(ctx context.Context, end time.Time) error {
	snap := p.db.NewSnapshot()
	defer snap.Close()

	var errs []error
	for _, ivl := range p.intervals {
		// Check if the given aggregation interval needs to be exported now
		if end.Truncate(ivl).Equal(end) {
			exportedCount, err := p.exportForInterval(ctx, snap, end, ivl)
			if err != nil {
				errs = append(errs, fmt.Errorf("failed to export interval %s for end time %d: %w", ivl, end.Unix(), err))
			}
			p.logger.Debug(
				"Finished exporting metrics",
				zap.Int("exported_datapoints", exportedCount),
				zap.Duration("interval", ivl),
				zap.Time("exported_till(exclusive)", end),
				zap.Error(err),
			)
		}
	}
	return errors.Join(errs...)
}

func (p *Processor) exportForInterval(
	ctx context.Context,
	snap *pebble.Snapshot,
	end time.Time,
	ivl time.Duration,
) (int, error) {
	from := merger.NewKey(ivl, zeroTime)
	lb, err := from.Marshal()
	if err != nil {
		return 0, fmt.Errorf("failed to encode range: %w", err)
	}

	to := merger.NewKey(ivl, end)
	ub, err := to.Marshal()
	if err != nil {
		return 0, fmt.Errorf("failed to encode range: %w", err)
	}

	iter, err := snap.NewIter(&pebble.IterOptions{
		LowerBound: lb,
		UpperBound: ub,
		KeyTypes:   pebble.IterKeyTypePointsOnly,
	})
	if err != nil {
		return 0, fmt.Errorf("failed to create iterator: %w", err)
	}
	defer iter.Close()

	var errs []error
	var exportedDPCount int
	for iter.First(); iter.Valid(); iter.Next() {
		var v merger.Value
		if err := v.UnmarshalProto(iter.Value()); err != nil {
			errs = append(errs, fmt.Errorf("failed to decode binary from database: %w", err))
			continue
		}
		if err := p.next.ConsumeMetrics(ctx, v.Metrics); err != nil {
			errs = append(errs, fmt.Errorf("failed to consume the decoded value: %w", err))
			continue
		}
		exportedDPCount += v.Metrics.DataPointCount()
	}
	if err := p.db.DeleteRange(lb, ub, p.wOpts); err != nil {
		errs = append(errs, fmt.Errorf("failed to delete exported entries: %w", err))
	}
	if len(errs) > 0 {
		return exportedDPCount, errors.Join(errs...)
	}
	return exportedDPCount, nil
}