builder.go

// Copyright © 2020 Jonathan Whitaker <github@whitaker.io>.
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.

package machine

import (
	"bytes"
	"context"
	"encoding/gob"
	"fmt"

	"github.com/google/uuid"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/trace"
)

// Stream is a representation of a data stream and its associated logic.
// It may be used individually or hosted by a Pipe. Creating a new Stream
// is handled by the appropriately named NewStream function.
//
// The Builder method is the entrypoint into creating the data processing flow.
// All branches of the Stream are required to end in either a Publish or
// a Link in order to be considered valid.
type Stream interface {
	ID() string
	Run(ctx context.Context, recorders ...recorder) error
	Inject(ctx context.Context, events map[string][]*Packet)
	Builder() Builder
}

// Builder is the interface provided for creating a data processing stream.
type Builder interface {
	Map(id string, a Applicative, options ...*Option) Builder
	FoldLeft(id string, f Fold, options ...*Option) Builder
	FoldRight(id string, f Fold, options ...*Option) Builder
	Fork(id string, f Fork, options ...*Option) (Builder, Builder)
	Loop(id string, x Fork, options ...*Option) (loop LoopBuilder, out Builder)
	Link(id, target string, options ...*Option)
	Publish(id string, s Publisher, options ...*Option)
}

// LoopBuilder is the interface provided for creating a data processing stream that loops.
type LoopBuilder interface {
	Map(id string, a Applicative, options ...*Option) LoopBuilder
	FoldLeft(id string, f Fold, options ...*Option) LoopBuilder
	FoldRight(id string, f Fold, options ...*Option) LoopBuilder
	Fork(id string, f Fork, options ...*Option) (LoopBuilder, LoopBuilder)
	Loop(id string, x Fork, options ...*Option) (loop LoopBuilder, out LoopBuilder)
	Publish(id string, s Publisher, options ...*Option)
	Done()
}

type nexter func(*node) *node

type looper struct {
	loopstart string
	next      Builder
}

type builder struct {
	vertex
	next      *node
	recorder  recorder
	vertacies map[string]*vertex
}

type node struct {
	vertex
	next  *node
	left  *node
	right *node
}

// ID is a method used to return the ID for the system
func (m *builder) ID() string {
	return m.id
}

// Run is the method used for starting the stream processing. It requires a context
// and an optional list of recorder functions. The recorder function has the signiture
// func(vertexID, vertexType, state string, paylaod []*Packet) and is called at the
// beginning of every vertex.
func (m *builder) Run(ctx context.Context, recorders ...recorder) error {
	if m.next == nil {
		return fmt.Errorf("non-terminated builder")
	}

	if len(recorders) > 0 {
		m.recorder = mergeRecorders(recorders...)
	}

	return m.cascade(ctx, m, m.input)
}

// Inject is a method for restarting work that has been dropped by the Stream
// typically in a distributed system setting. Though it can be used to side load
// data into the Stream to be processed
func (m *builder) Inject(ctx context.Context, events map[string][]*Packet) {
	for node, payload := range events {
		if v, ok := m.vertacies[node]; ok {
			if !*v.option.Injectable {
				m.recorder(v.id, v.vertexType, "injection-denied", payload)
				continue
			}

			v.input.channel <- payload
		}
	}
}

func (m *builder) Builder() Builder {
	return nexter(func(n *node) *node {
		m.next = n
		return n
	})
}

// Map apply a mutation, options default to the set used when creating the Stream
func (n nexter) Map(id string, x Applicative, options ...*Option) Builder {
	opt := &Option{
		BufferSize: intP(0),
	}

	if len(options) > 0 {
		opt = opt.merge(options...)
	}

	next := &node{}
	edge := newEdge(opt.BufferSize)

	next.vertex = vertex{
		id:         id,
		vertexType: "map",
		option:     opt,
		handler: func(payload []*Packet) {
			for _, packet := range payload {
				packet.apply(id, x)
			}

			edge.channel <- payload
		},
		connector: func(ctx context.Context, b *builder) error {
			if next.next == nil {
				return fmt.Errorf("non-terminated map")
			}
			return next.next.cascade(ctx, b, edge)
		},
	}

	next = n(next)

	return nexter(func(n *node) *node {
		next.next = n
		return n
	})
}

// FoldLeft the data, options default to the set used when creating the Stream
func (n nexter) FoldLeft(id string, x Fold, options ...*Option) Builder {
	opt := &Option{
		BufferSize: intP(0),
	}

	if len(options) > 0 {
		opt = opt.merge(options...)
	}

	next := &node{}
	edge := newEdge(opt.BufferSize)

	fr := func(payload ...*Packet) *Packet {
		if len(payload) == 1 {
			return payload[0]
		}

		d := payload[0]

		for i := 1; i < len(payload); i++ {
			d.Data = x(d.Data, payload[i].Data)
		}

		return d
	}

	next.vertex = vertex{
		id:         id,
		vertexType: "fold",
		option:     opt,
		handler: func(payload []*Packet) {
			edge.channel <- []*Packet{fr(payload...)}
		},
		connector: func(ctx context.Context, b *builder) error {
			if next.next == nil {
				return fmt.Errorf("non-terminated fold")
			}
			return next.next.cascade(ctx, b, edge)
		},
	}
	next = n(next)

	return nexter(func(n *node) *node {
		next.next = n
		return n
	})
}

// FoldRight the data, options default to the set used when creating the Stream
func (n nexter) FoldRight(id string, x Fold, options ...*Option) Builder {
	opt := &Option{
		BufferSize: intP(0),
	}

	if len(options) > 0 {
		opt = opt.merge(options...)
	}

	next := &node{}
	edge := newEdge(opt.BufferSize)

	var fr func(...*Packet) *Packet
	fr = func(payload ...*Packet) *Packet {
		if len(payload) == 1 {
			return payload[0]
		}

		payload[len(payload)-1].Data = x(payload[0].Data, fr(payload[1:]...).Data)

		return payload[len(payload)-1]
	}

	next.vertex = vertex{
		id:         id,
		vertexType: "fold",
		option:     opt,
		handler: func(payload []*Packet) {
			edge.channel <- []*Packet{fr(payload...)}
		},
		connector: func(ctx context.Context, b *builder) error {
			if next.next == nil {
				return fmt.Errorf("non-terminated node")
			}
			return next.next.cascade(ctx, b, edge)
		},
	}

	next = n(next)

	return nexter(func(n *node) *node {
		next.next = n
		return n
	})
}

// Fork the data, options default to the set used when creating the Stream
func (n nexter) Fork(id string, x Fork, options ...*Option) (left, right Builder) {
	opt := &Option{
		BufferSize: intP(0),
	}

	if len(options) > 0 {
		opt = opt.merge(options...)
	}

	next := &node{}

	leftEdge := newEdge(opt.BufferSize)
	rightEdge := newEdge(opt.BufferSize)

	next.vertex = vertex{
		id:         id,
		vertexType: "fork",
		option:     opt,
		handler: func(payload []*Packet) {
			lpayload, rpayload := x(payload)
			leftEdge.channel <- lpayload
			rightEdge.channel <- rpayload
		},
		connector: func(ctx context.Context, b *builder) error {
			if next.left == nil || next.right == nil {
				return fmt.Errorf("non-terminated fork")
			} else if err := next.left.cascade(ctx, b, leftEdge); err != nil {
				return err
			} else if err := next.right.cascade(ctx, b, rightEdge); err != nil {
				return err
			}

			return nil
		},
	}

	next = n(next)

	return nexter(func(n *node) *node {
			next.left = n
			return n
		}), nexter(func(n *node) *node {
			next.right = n
			return n
		})
}

// Link the data to an existing operation creating a loop,
// target must be an ancestor, options default to the set
// used when creating the Stream
func (n nexter) Link(id, target string, options ...*Option) {
	opt := &Option{
		BufferSize: intP(0),
	}

	if len(options) > 0 {
		opt = opt.merge(options...)
	}

	edge := newEdge(opt.BufferSize)
	n(&node{
		vertex: vertex{
			id:         id,
			vertexType: "link",
			handler:    func(payload []*Packet) { edge.channel <- payload },
			option:     opt,
			connector: func(ctx context.Context, b *builder) error {
				v, ok := b.vertacies[target]

				if !ok {
					return fmt.Errorf("invalid target - not in list of ancestors")
				}

				return v.cascade(ctx, b, edge)
			},
		},
	})
}

// Publish the data outside the system, options default to the set used when creating the Stream
func (n nexter) Publish(id string, x Publisher, options ...*Option) {
	opt := &Option{
		BufferSize: intP(0),
	}

	if len(options) > 0 {
		opt = opt.merge(options...)
	}

	n(&node{
		vertex: vertex{
			id:         id,
			vertexType: "publish",
			option:     opt,
			handler: func(payload []*Packet) {
				data := make([]Data, len(payload))
				for i, packet := range payload {
					data[i] = packet.Data
				}

				if err := x.Send(data); err != nil {
					for _, packet := range payload {
						packet.handleError(id, err)
					}
				}
			},
			connector: func(ctx context.Context, b *builder) error { return nil },
		},
	})
}

// Loop the data combining a fork and link the first output is the Builder for the loop
// and the second is the output of the loop
func (n nexter) Loop(id string, x Fork, options ...*Option) (loop LoopBuilder, out Builder) {
	left, right := n.Fork(id, x, options...)

	return &looper{
			loopstart: id,
			next:      left,
		},
		right
}

// Map apply a mutation, options default to the set used when creating the Stream
func (n *looper) Map(id string, a Applicative, options ...*Option) LoopBuilder {
	return &looper{
		loopstart: n.loopstart,
		next:      n.next.Map(id, a, options...),
	}
}

// FoldLeft the data, options default to the set used when creating the Stream
func (n *looper) FoldLeft(id string, f Fold, options ...*Option) LoopBuilder {
	return &looper{
		loopstart: n.loopstart,
		next:      n.next.FoldLeft(id, f, options...),
	}
}

// FoldRight the data, options default to the set used when creating the Stream
func (n *looper) FoldRight(id string, f Fold, options ...*Option) LoopBuilder {
	return &looper{
		loopstart: n.loopstart,
		next:      n.next.FoldRight(id, f, options...),
	}
}

// Fork the data, options default to the set used when creating the Stream
func (n *looper) Fork(id string, f Fork, options ...*Option) (left, right LoopBuilder) {
	x, y := n.next.Fork(id, f, options...)

	return &looper{
			loopstart: n.loopstart,
			next:      x,
		}, &looper{
			loopstart: n.loopstart,
			next:      y,
		}
}

// Loop the data combining a fork and link the first output is the Builder for the loop
// and the second is the output of the loop
func (n *looper) Loop(id string, x Fork, options ...*Option) (loop, out LoopBuilder) {
	inner, outer := n.next.Loop(id, x, options...)
	return inner, &looper{
		loopstart: n.loopstart,
		next:      outer,
	}
}

// Publish the data outside the system, options default to the set used when creating the Stream
func (n *looper) Publish(id string, x Publisher, options ...*Option) {
	n.next.Publish(id, x, options...)
}

// Done function for ending the loop and going back to the original Fork
func (n *looper) Done() {
	n.next.Link(uuid.NewString(), n.loopstart)
}

// NewStream is a function for creating a new Stream. It takes an id, a Retriever function,
// and a list of Options that can override the defaults and set new defaults for the
// subsequent vertices in the Stream.
func NewStream(id string, retriever Retriever, options ...*Option) Stream {
	opt := defaultOptions.merge(options...)

	tracer := otel.GetTracerProvider().Tracer("stream" + "." + id)
	vertexAttributes := trace.WithAttributes(
		attribute.String("vertex_id", id),
		attribute.String("vertex_type", "stream"),
	)

	edge := newEdge(opt.BufferSize)
	input := newEdge(opt.BufferSize)

	x := &builder{
		vertex: vertex{
			id:         id,
			vertexType: "stream",
			input:      input,
			option:     opt,
			handler: func(p []*Packet) {
				edge.channel <- p
			},
		},
		vertacies: map[string]*vertex{},
		recorder:  func(s1, s2, s3 string, p []*Packet) {},
	}

	x.connector = func(ctx context.Context, b *builder) error {
		i := retriever(ctx)

		go func() {
		Loop:
			for {
				select {
				case <-ctx.Done():
					break Loop
				case data := <-i:
					if len(data) < 1 {
						continue
					}

					payload := make([]*Packet, len(data))
					for i, item := range data {
						var id string

						if val, ok := item["__traceID"]; ok && *opt.TraceID {
							if strval, ok := val.(string); ok {
								id = strval
							}
						}

						if id == "" {
							id = uuid.NewString()
						}

						packet := &Packet{
							ID:   id,
							Data: item,
						}
						if *x.option.Span {
							packet.newSpan(ctx, tracer, "stream.begin", vertexAttributes)
						}
						payload[i] = packet
					}

					input.channel <- payload
				}
			}
		}()
		return x.next.cascade(ctx, x, edge)
	}

	return x
}

func mergeRecorders(recorders ...recorder) recorder {
	return func(id, name string, state string, payload []*Packet) {
		out := []*Packet{}
		buf := &bytes.Buffer{}
		enc, dec := gob.NewEncoder(buf), gob.NewDecoder(buf)

		_ = enc.Encode(payload)
		_ = dec.Decode(&out)

		for _, r := range recorders {
			r(id, name, state, out)
		}
	}
}

func init() {
	gob.Register([]*Packet{})
	gob.Register([]Data{})
}