collector.go

// Package data contains the code involved with Viam's Data Management Platform for automatically collecting component
// readings from robots.
package data

import (
	"context"
	"fmt"
	"reflect"
	"sync"
	"sync/atomic"
	"time"

	"github.com/benbjohnson/clock"
	"github.com/pkg/errors"
	"go.opencensus.io/trace"
	v1 "go.viam.com/api/app/datasync/v1"
	pb "go.viam.com/api/common/v1"
	"go.viam.com/utils"
	"go.viam.com/utils/protoutils"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/types/known/anypb"
	"google.golang.org/protobuf/types/known/structpb"
	"google.golang.org/protobuf/types/known/timestamppb"

	"go.viam.com/rdk/logging"
	"go.viam.com/rdk/resource"
	"go.viam.com/rdk/services/datamanager/datacapture"
)

// The cutoff at which if interval < cutoff, a sleep based capture func is used instead of a ticker.
var sleepCaptureCutoff = 2 * time.Millisecond

// CaptureFunc allows the creation of simple Capturers with anonymous functions.
type CaptureFunc func(ctx context.Context, params map[string]*anypb.Any) (interface{}, error)

// FromDMContextKey is used to check whether the context is from data management.
// Deprecated: use a camera.Extra with camera.NewContext instead.
type FromDMContextKey struct{}

// FromDMString is used to access the 'fromDataManagement' value from a request's Extra struct.
const FromDMString = "fromDataManagement"

// FromDMExtraMap is a map with 'fromDataManagement' set to true.
var FromDMExtraMap = map[string]interface{}{FromDMString: true}

// ErrNoCaptureToStore is returned when a modular filter resource filters the capture coming from the base resource.
var ErrNoCaptureToStore = status.Error(codes.FailedPrecondition, "no capture from filter module")

// If an error is ongoing, the frequency (in seconds) with which to suppress identical error logs.
const identicalErrorLogFrequencyHz = 2

// Collector collects data to some target.
type Collector interface {
	Close()
	Collect()
	Flush()
}

type collector struct {
	clock          clock.Clock
	captureResults chan *v1.SensorData
	captureErrors  chan error
	interval       time.Duration
	params         map[string]*anypb.Any
	// `lock` serializes calls to `Flush` and `Close`.
	lock             sync.Mutex
	logger           logging.Logger
	captureWorkers   sync.WaitGroup
	logRoutine       sync.WaitGroup
	cancelCtx        context.Context
	cancel           context.CancelFunc
	captureFunc      CaptureFunc
	closeStarted     atomic.Bool
	closeFinished    bool
	target           datacapture.BufferedWriter
	lastLoggedErrors map[string]int64
}

// Close closes the channels backing the Collector. It should always be called before disposing of a Collector to avoid
// leaking goroutines.
func (c *collector) Close() {
	// Signal the start of closing. Specifically, this is used to signal to the error logging
	// goroutine that it should ignore "context canceled" errors.
	c.closeStarted.Store(true)

	// Signal all `captureWorkers` to exit.
	c.cancel()
	// CaptureWorkers acquire the `c.lock` to do their work (i.e: call `collector.Flush()`). We must
	// `Wait` on them before acquiring the lock to avoid deadlock.
	c.captureWorkers.Wait()

	c.lock.Lock()
	defer c.lock.Unlock()
	if c.closeFinished {
		return
	}

	if err := c.target.Flush(); err != nil {
		c.logger.Errorw("failed to flush capture data", "error", err)
	}

	close(c.captureErrors)
	c.logRoutine.Wait()

	// Closing has completed. Set this to prevent a double-close from executing.
	c.closeFinished = true
}

func (c *collector) Flush() {
	c.lock.Lock()
	defer c.lock.Unlock()
	if err := c.target.Flush(); err != nil {
		c.logger.Errorw("failed to flush collector", "error", err)
	}
}

// Collect starts the Collector, causing it to run c.capturer.Capture every c.interval, and write the results to
// c.target. It blocks until the underlying capture goroutine starts.
func (c *collector) Collect() {
	_, span := trace.StartSpan(c.cancelCtx, "data::collector::Collect")
	defer span.End()

	started := make(chan struct{})
	c.captureWorkers.Add(1)
	utils.PanicCapturingGo(func() {
		defer c.captureWorkers.Done()
		c.capture(started)
	})
	c.captureWorkers.Add(1)
	utils.PanicCapturingGo(func() {
		defer c.captureWorkers.Done()
		if err := c.writeCaptureResults(); err != nil {
			c.captureErrors <- errors.Wrap(err, fmt.Sprintf("failed to write to collector %s", c.target.Path()))
		}
	})
	c.logRoutine.Add(1)
	utils.PanicCapturingGo(func() {
		defer c.logRoutine.Done()
		c.logCaptureErrs()
	})
	<-started
}

// Go's time.Ticker has inconsistent performance with durations of below 1ms [0], so we use a time.Sleep based approach
// when the configured capture interval is below 2ms. A Ticker based approach is kept for longer capture intervals to
// avoid wasting CPU on a thread that's idling for the vast majority of the time.
// [0]: https://www.mail-archive.com/golang-nuts@googlegroups.com/msg46002.html
func (c *collector) capture(started chan struct{}) {
	if c.interval < sleepCaptureCutoff {
		c.sleepBasedCapture(started)
	} else {
		c.tickerBasedCapture(started)
	}
}

func (c *collector) sleepBasedCapture(started chan struct{}) {
	next := c.clock.Now().Add(c.interval)
	until := c.clock.Until(next)
	var captureWorkers sync.WaitGroup

	close(started)
	for {
		if err := c.cancelCtx.Err(); err != nil {
			c.captureErrors <- errors.Wrap(err, "error in context")
			captureWorkers.Wait()
			close(c.captureResults)
			return
		}
		c.clock.Sleep(until)

		select {
		case <-c.cancelCtx.Done():
			captureWorkers.Wait()
			close(c.captureResults)
			return
		default:
			captureWorkers.Add(1)
			utils.PanicCapturingGo(func() {
				defer captureWorkers.Done()
				c.getAndPushNextReading()
			})
		}
		next = next.Add(c.interval)
		until = c.clock.Until(next)
	}
}

func (c *collector) tickerBasedCapture(started chan struct{}) {
	ticker := c.clock.Ticker(c.interval)
	defer ticker.Stop()
	var captureWorkers sync.WaitGroup

	close(started)
	for {
		if err := c.cancelCtx.Err(); err != nil {
			c.captureErrors <- errors.Wrap(err, "error in context")
			captureWorkers.Wait()
			close(c.captureResults)
			return
		}
		select {
		case <-c.cancelCtx.Done():
			captureWorkers.Wait()
			close(c.captureResults)
			return
		case <-ticker.C:
			captureWorkers.Add(1)
			utils.PanicCapturingGo(func() {
				defer captureWorkers.Done()
				c.getAndPushNextReading()
			})
		}
	}
}

func (c *collector) getAndPushNextReading() {
	timeRequested := timestamppb.New(c.clock.Now().UTC())
	reading, err := c.captureFunc(c.cancelCtx, c.params)
	timeReceived := timestamppb.New(c.clock.Now().UTC())
	if err != nil {
		if errors.Is(err, ErrNoCaptureToStore) {
			c.logger.Debug("capture filtered out by modular resource")
			return
		}
		c.captureErrors <- errors.Wrap(err, "error while capturing data")
		return
	}

	var msg v1.SensorData
	switch v := reading.(type) {
	case []byte:
		msg = v1.SensorData{
			Metadata: &v1.SensorMetadata{
				TimeRequested: timeRequested,
				TimeReceived:  timeReceived,
			},
			Data: &v1.SensorData_Binary{
				Binary: v,
			},
		}
	default:
		// If it's not bytes, it's a struct.
		var pbReading *structpb.Struct
		var err error

		if reflect.TypeOf(reading) == reflect.TypeOf(pb.GetReadingsResponse{}) {
			// We special-case the GetReadingsResponse because it already contains
			// structpb.Values in it, and the StructToStructPb logic does not handle
			// that cleanly.
			topLevelMap := make(map[string]*structpb.Value)
			topLevelMap["readings"] = structpb.NewStructValue(
				&structpb.Struct{Fields: reading.(pb.GetReadingsResponse).Readings},
			)
			pbReading = &structpb.Struct{Fields: topLevelMap}
		} else {
			pbReading, err = protoutils.StructToStructPbIgnoreOmitEmpty(reading)
			if err != nil {
				c.captureErrors <- errors.Wrap(err, "error while converting reading to structpb.Struct")
				return
			}
		}

		msg = v1.SensorData{
			Metadata: &v1.SensorMetadata{
				TimeRequested: timeRequested,
				TimeReceived:  timeReceived,
			},
			Data: &v1.SensorData_Struct{
				Struct: pbReading,
			},
		}
	}

	select {
	// If c.captureResults is full, c.captureResults <- a can block indefinitely. This additional select block allows cancel to
	// still work when this happens.
	case <-c.cancelCtx.Done():
	case c.captureResults <- &msg:
	}
}

// NewCollector returns a new Collector with the passed capturer and configuration options. It calls capturer at the
// specified Interval, and appends the resulting reading to target.
func NewCollector(captureFunc CaptureFunc, params CollectorParams) (Collector, error) {
	if err := params.Validate(); err != nil {
		return nil, errors.Wrap(err, fmt.Sprintf("failed to construct collector for %s", params.ComponentName))
	}

	cancelCtx, cancelFunc := context.WithCancel(context.Background())
	var c clock.Clock
	if params.Clock == nil {
		c = clock.New()
	} else {
		c = params.Clock
	}
	return &collector{
		captureResults:   make(chan *v1.SensorData, params.QueueSize),
		captureErrors:    make(chan error, params.QueueSize),
		interval:         params.Interval,
		params:           params.MethodParams,
		logger:           params.Logger,
		cancelCtx:        cancelCtx,
		cancel:           cancelFunc,
		captureFunc:      captureFunc,
		target:           params.Target,
		clock:            c,
		lastLoggedErrors: make(map[string]int64, 0),
	}, nil
}

func (c *collector) writeCaptureResults() error {
	for msg := range c.captureResults {
		if err := c.target.Write(msg); err != nil {
			return err
		}
	}
	return nil
}

func (c *collector) logCaptureErrs() {
	for err := range c.captureErrors {
		now := c.clock.Now().Unix()
		if c.closeStarted.Load() {
			// Don't log context cancellation errors if the collector has already been closed. This
			// means the collector canceled the context, and the context cancellation error is
			// expected.
			if errors.Is(err, context.Canceled) {
				continue
			}
		}
		// Only log a specific error message if we haven't logged it in the past 2 seconds.
		if lastLogged, ok := c.lastLoggedErrors[err.Error()]; (ok && int(now-lastLogged) > identicalErrorLogFrequencyHz) || !ok {
			c.logger.Error((err))
			c.lastLoggedErrors[err.Error()] = now
		}
	}
}

// InvalidInterfaceErr is the error describing when an interface not conforming to the expected resource.API was
// passed into a CollectorConstructor.
func InvalidInterfaceErr(api resource.API) error {
	return errors.Errorf("passed interface does not conform to expected resource type %s", api)
}

// FailedToReadErr is the error describing when a Capturer was unable to get the reading of a method.
func FailedToReadErr(component, method string, err error) error {
	return errors.Errorf("failed to get reading of method %s of component %s: %v", method, component, err)
}

// GetExtraFromContext sets the extra struct with "fromDataManagement": true if the flag is true in the context.
// Deprecated: Use camera.FromContext instead.
func GetExtraFromContext(ctx context.Context) (*structpb.Struct, error) {
	extra := make(map[string]interface{})
	if ctx.Value(FromDMContextKey{}) == true {
		extra[FromDMString] = true
	}
	return protoutils.StructToStructPb(extra)
}