diff --git a/cmd/commands/processor.go b/cmd/commands/processor.go index 520759bb59..b394603d45 100644 --- a/cmd/commands/processor.go +++ b/cmd/commands/processor.go @@ -28,6 +28,7 @@ import ( "github.com/numaproj/numaflow" dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1" + "github.com/numaproj/numaflow/pkg/flatmap" "github.com/numaproj/numaflow/pkg/shared/logging" "github.com/numaproj/numaflow/pkg/sinks" "github.com/numaproj/numaflow/pkg/sources" @@ -95,6 +96,18 @@ func NewProcessorCommand() *cobra.Command { ISBSvcType: dfv1.ISBSvcType(isbSvcType), VertexInstance: vertexInstance, } + + enableFlatMapUdfStream, err := vertexInstance.Vertex.FlatmapUdfStreamEnabled() + if err != nil { + return fmt.Errorf("failed to parse Flatmap streaming UDF metadata, %w", err) + } + if enableFlatMapUdfStream { + flatMapProcessor := &flatmap.FlatmapUDFProcessor{ + ISBSvcType: dfv1.ISBSvcType(isbSvcType), + VertexInstance: vertexInstance, + } + return flatMapProcessor.Start(ctx) + } return p.Start(ctx) case dfv1.VertexTypeReduceUDF: p := &udf.ReduceUDFProcessor{ diff --git a/pkg/apis/numaflow/v1alpha1/const.go b/pkg/apis/numaflow/v1alpha1/const.go index f129efdf53..dfd002ee1d 100644 --- a/pkg/apis/numaflow/v1alpha1/const.go +++ b/pkg/apis/numaflow/v1alpha1/const.go @@ -200,6 +200,9 @@ const ( // UDF map streaming MapUdfStreamKey = "numaflow.numaproj.io/map-stream" + // FlatmapUdfStreamKey is used to annotate a vertex to use the + // bidirectional streaming gRPC for flatmap + FlatmapUdfStreamKey = "numaflow.numaproj.io/flat-map" // Pipeline health status PipelineStatusHealthy = "healthy" diff --git a/pkg/apis/numaflow/v1alpha1/vertex_types.go b/pkg/apis/numaflow/v1alpha1/vertex_types.go index 74663cec72..a5d8c5a0e2 100644 --- a/pkg/apis/numaflow/v1alpha1/vertex_types.go +++ b/pkg/apis/numaflow/v1alpha1/vertex_types.go @@ -421,6 +421,15 @@ func (v Vertex) MapUdfStreamEnabled() (bool, error) { return false, nil } +func (v Vertex) FlatmapUdfStreamEnabled() (bool, error) { + if v.Spec.Metadata != nil && v.Spec.Metadata.Annotations != nil { + if flatmapUdfStream, existing := v.Spec.Metadata.Annotations[FlatmapUdfStreamKey]; existing { + return strconv.ParseBool(flatmapUdfStream) + } + } + return false, nil +} + type VertexSpec struct { AbstractVertex `json:",inline" protobuf:"bytes,1,opt,name=abstractVertex"` PipelineName string `json:"pipelineName" protobuf:"bytes,2,opt,name=pipelineName"` diff --git a/pkg/flatmap/common.go b/pkg/flatmap/common.go new file mode 100644 index 0000000000..6586bf8ccc --- /dev/null +++ b/pkg/flatmap/common.go @@ -0,0 +1,171 @@ +/* +Copyright 2022 The Numaproj Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package flatmap + +import ( + "context" + "fmt" + "strings" + + dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1" + "github.com/numaproj/numaflow/pkg/isb" + jetstreamisb "github.com/numaproj/numaflow/pkg/isb/stores/jetstream" + redisisb "github.com/numaproj/numaflow/pkg/isb/stores/redis" + "github.com/numaproj/numaflow/pkg/isbsvc" + jsclient "github.com/numaproj/numaflow/pkg/shared/clients/nats" + redisclient "github.com/numaproj/numaflow/pkg/shared/clients/redis" +) + +func buildRedisBufferIO(ctx context.Context, vertexInstance *dfv1.VertexInstance) ([]isb.BufferReader, map[string][]isb.BufferWriter, error) { + var readers []isb.BufferReader + redisClient := redisclient.NewInClusterRedisClient() + var readerOpts []redisclient.Option + if x := vertexInstance.Vertex.Spec.Limits; x != nil && x.ReadTimeout != nil { + readerOpts = append(readerOpts, redisclient.WithReadTimeOut(x.ReadTimeout.Duration)) + } + // create readers for owned buffer partitions. + // For reduce vertex, we only need to read from one buffer partition. + if vertexInstance.Vertex.GetVertexType() == dfv1.VertexTypeReduceUDF { + var fromBufferPartition string + // find the buffer partition owned by this replica. + for _, b := range vertexInstance.Vertex.OwnedBuffers() { + if strings.HasSuffix(b, fmt.Sprintf("-%d", vertexInstance.Replica)) { + fromBufferPartition = b + break + } + } + if len(fromBufferPartition) == 0 { + return nil, nil, fmt.Errorf("can not find from buffer") + } + + fromGroup := fromBufferPartition + "-group" + consumer := fmt.Sprintf("%s-%v", vertexInstance.Vertex.Name, vertexInstance.Replica) + // since we read from one buffer partition, fromPartitionIdx is 0. + reader := redisisb.NewBufferRead(ctx, redisClient, fromBufferPartition, fromGroup, consumer, 0, readerOpts...) + readers = append(readers, reader) + } else { + // for map vertex, we need to read from all buffer partitions. So create readers for all buffer partitions. + for _, bufferPartition := range vertexInstance.Vertex.OwnedBuffers() { + fromGroup := bufferPartition + "-group" + consumer := fmt.Sprintf("%s-%v", vertexInstance.Vertex.Name, vertexInstance.Replica) + reader := redisisb.NewBufferRead(ctx, redisClient, bufferPartition, fromGroup, consumer, 0, readerOpts...) + readers = append(readers, reader) + } + } + + // create writers for toVertex's buffer partitions. + // we create a map of toVertex -> []BufferWriter(writer for each partition) + writers := make(map[string][]isb.BufferWriter) + for _, e := range vertexInstance.Vertex.Spec.ToEdges { + + writeOpts := []redisclient.Option{ + redisclient.WithBufferFullWritingStrategy(e.BufferFullWritingStrategy()), + } + if x := e.ToVertexLimits; x != nil && x.BufferMaxLength != nil { + writeOpts = append(writeOpts, redisclient.WithMaxLength(int64(*x.BufferMaxLength))) + } + if x := e.ToVertexLimits; x != nil && x.BufferUsageLimit != nil { + writeOpts = append(writeOpts, redisclient.WithBufferUsageLimit(float64(*x.BufferUsageLimit)/100)) + } + var edgeBuffers []isb.BufferWriter + partitionedBuffers := dfv1.GenerateBufferNames(vertexInstance.Vertex.Namespace, vertexInstance.Vertex.Spec.PipelineName, e.To, e.GetToVertexPartitionCount()) + for partitionIdx, partition := range partitionedBuffers { + writer := redisisb.NewBufferWrite(ctx, redisClient, partition, partition+"-group", int32(partitionIdx), writeOpts...) + edgeBuffers = append(edgeBuffers, writer) + } + writers[e.To] = edgeBuffers + } + + return readers, writers, nil +} + +func buildJetStreamBufferIO(ctx context.Context, vertexInstance *dfv1.VertexInstance, clientPool *jsclient.ClientPool) ([]isb.BufferReader, map[string][]isb.BufferWriter, error) { + + // create readers for owned buffer partitions. + var readers []isb.BufferReader + var readOptions []jetstreamisb.ReadOption + if x := vertexInstance.Vertex.Spec.Limits; x != nil && x.ReadTimeout != nil { + readOptions = append(readOptions, jetstreamisb.WithReadTimeOut(x.ReadTimeout.Duration)) + } + + // create readers for owned buffer partitions. + // For reduce vertex, we only need to read from one buffer partition. + if vertexInstance.Vertex.GetVertexType() == dfv1.VertexTypeReduceUDF { + // choose the buffer that corresponds to this reduce processor because + // reducer's incoming buffer can have more than one partition for parallelism + var fromBufferPartition string + // find the buffer partition owned by this replica. + for _, b := range vertexInstance.Vertex.OwnedBuffers() { + if strings.HasSuffix(b, fmt.Sprintf("-%d", vertexInstance.Replica)) { + fromBufferPartition = b + break + } + } + if len(fromBufferPartition) == 0 { + return nil, nil, fmt.Errorf("can not find from buffer") + } + + fromStreamName := isbsvc.JetStreamName(fromBufferPartition) + // reduce processor only has one buffer partition + // since we read from one buffer partition, fromPartitionIdx is 0. + reader, err := jetstreamisb.NewJetStreamBufferReader(ctx, clientPool.NextAvailableClient(), fromBufferPartition, fromStreamName, fromStreamName, 0, readOptions...) + if err != nil { + return nil, nil, err + } + readers = append(readers, reader) + } else { + // for map vertex, we need to read from all buffer partitions. So create readers for all buffer partitions. + for index, bufferPartition := range vertexInstance.Vertex.OwnedBuffers() { + fromStreamName := isbsvc.JetStreamName(bufferPartition) + + reader, err := jetstreamisb.NewJetStreamBufferReader(ctx, clientPool.NextAvailableClient(), bufferPartition, fromStreamName, fromStreamName, int32(index), readOptions...) + if err != nil { + return nil, nil, err + } + readers = append(readers, reader) + } + } + + // create writers for toVertex's buffer partitions. + // we create a map of toVertex -> []BufferWriter(writer for each partition) + writers := make(map[string][]isb.BufferWriter) + for _, e := range vertexInstance.Vertex.Spec.ToEdges { + writeOpts := []jetstreamisb.WriteOption{ + jetstreamisb.WithBufferFullWritingStrategy(e.BufferFullWritingStrategy()), + } + if x := e.ToVertexLimits; x != nil && x.BufferMaxLength != nil { + writeOpts = append(writeOpts, jetstreamisb.WithMaxLength(int64(*x.BufferMaxLength))) + } + if x := e.ToVertexLimits; x != nil && x.BufferUsageLimit != nil { + writeOpts = append(writeOpts, jetstreamisb.WithBufferUsageLimit(float64(*x.BufferUsageLimit)/100)) + } + + partitionedBuffers := dfv1.GenerateBufferNames(vertexInstance.Vertex.Namespace, vertexInstance.Vertex.Spec.PipelineName, e.To, e.GetToVertexPartitionCount()) + var edgeBuffers []isb.BufferWriter + for partitionIdx, partition := range partitionedBuffers { + streamName := isbsvc.JetStreamName(partition) + writer, err := jetstreamisb.NewJetStreamBufferWriter(ctx, clientPool.NextAvailableClient(), partition, streamName, streamName, int32(partitionIdx), writeOpts...) + if err != nil { + return nil, nil, err + } + edgeBuffers = append(edgeBuffers, writer) + } + + writers[e.To] = edgeBuffers + } + return readers, writers, nil +} diff --git a/pkg/flatmap/flatmap_udf.go b/pkg/flatmap/flatmap_udf.go new file mode 100644 index 0000000000..146f2d6772 --- /dev/null +++ b/pkg/flatmap/flatmap_udf.go @@ -0,0 +1,285 @@ +/* +Copyright 2022 The Numaproj Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package flatmap + +import ( + "context" + "fmt" + "sync" + + "go.uber.org/zap" + + dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1" + "github.com/numaproj/numaflow/pkg/flatmap/forward" + "github.com/numaproj/numaflow/pkg/flatmap/rpc" + "github.com/numaproj/numaflow/pkg/forwarder" + "github.com/numaproj/numaflow/pkg/isb" + "github.com/numaproj/numaflow/pkg/metrics" + "github.com/numaproj/numaflow/pkg/sdkclient" + "github.com/numaproj/numaflow/pkg/sdkclient/flatmapper" + sdkserverinfo "github.com/numaproj/numaflow/pkg/sdkclient/serverinfo" + jsclient "github.com/numaproj/numaflow/pkg/shared/clients/nats" + "github.com/numaproj/numaflow/pkg/shared/logging" + sharedutil "github.com/numaproj/numaflow/pkg/shared/util" + "github.com/numaproj/numaflow/pkg/shuffle" + "github.com/numaproj/numaflow/pkg/watermark/fetch" + "github.com/numaproj/numaflow/pkg/watermark/generic" + "github.com/numaproj/numaflow/pkg/watermark/generic/jetstream" + "github.com/numaproj/numaflow/pkg/watermark/store" + "github.com/numaproj/numaflow/pkg/watermark/wmb" +) + +type FlatmapUDFProcessor struct { + ISBSvcType dfv1.ISBSvcType + VertexInstance *dfv1.VertexInstance +} + +func (u *FlatmapUDFProcessor) Start(ctx context.Context) error { + log := logging.FromContext(ctx) + finalWg := sync.WaitGroup{} + + ctx, cancel := context.WithCancel(ctx) + defer cancel() + + fromBuffer := u.VertexInstance.Vertex.OwnedBuffers() + log = log.With("protocol", "uds-grpc-flatmap-udf") + + var ( + readers []isb.BufferReader + writers map[string][]isb.BufferWriter + fromVertexWmStores map[string]store.WatermarkStore + toVertexWmStores map[string]store.WatermarkStore + flatmapHandler *rpc.GRPCBasedFlatmap + idleManager wmb.IdleManager + ) + + // watermark variables + fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferList(u.VertexInstance.Vertex.GetToBuffers()) + idleManager = wmb.NewNoOpIdleManager() + + var err error + + // create readers and writers + switch u.ISBSvcType { + case dfv1.ISBSvcTypeRedis: + readers, writers, err = buildRedisBufferIO(ctx, u.VertexInstance) + if err != nil { + return err + } + case dfv1.ISBSvcTypeJetStream: + + natsClientPool, err := jsclient.NewClientPool(ctx) + if err != nil { + return fmt.Errorf("failed to create a new NATS client pool: %w", err) + } + defer natsClientPool.CloseAll() + + // multiple go routines can share the same set of writers since nats conn is thread safe + // https://github.com/nats-io/nats.go/issues/241 + readers, writers, err = buildJetStreamBufferIO(ctx, u.VertexInstance, natsClientPool) + if err != nil { + return err + } + + // created watermark related components only if watermark is enabled + // otherwise no op will used + if !u.VertexInstance.Vertex.Spec.Watermark.Disabled { + // create from vertex watermark stores + fromVertexWmStores, err = jetstream.BuildFromVertexWatermarkStores(ctx, u.VertexInstance, natsClientPool.NextAvailableClient()) + if err != nil { + return fmt.Errorf("failed to build watermark stores: %w", err) + } + + // create watermark fetcher using watermark stores + fetchWatermark = fetch.NewEdgeFetcherSet(ctx, u.VertexInstance, fromVertexWmStores, fetch.WithVertexReplica(u.VertexInstance.Replica), + fetch.WithIsReduce(u.VertexInstance.Vertex.IsReduceUDF()), fetch.WithIsSource(u.VertexInstance.Vertex.IsASource())) + + // create to vertex watermark stores + toVertexWmStores, err = jetstream.BuildToVertexWatermarkStores(ctx, u.VertexInstance, natsClientPool.NextAvailableClient()) + if err != nil { + return err + } + + // create watermark publisher using watermark stores + publishWatermark = jetstream.BuildPublishersFromStores(ctx, u.VertexInstance, toVertexWmStores) + + idleManager, _ = wmb.NewIdleManager(len(writers), len(writers)) + } + default: + return fmt.Errorf("unrecognized isbsvc type %q", u.ISBSvcType) + } + + maxMessageSize := sharedutil.LookupEnvIntOr(dfv1.EnvGRPCMaxMessageSize, sdkclient.DefaultGRPCMaxMessageSize) + // Wait for server info to be ready + serverInfo, err := sdkserverinfo.SDKServerInfo(sdkserverinfo.WithServerInfoFilePath(sdkclient.FlatmapServerInfoFile)) + if err != nil { + return err + } + + flatmapClient, err := flatmapper.New(serverInfo, sdkclient.WithMaxMessageSize(maxMessageSize)) + if err != nil { + return fmt.Errorf("failed to create map client, %w", err) + } + flatmapHandler = rpc.NewUDSgRPCBasedFlatmap(flatmapClient, u.VertexInstance.Vertex.Spec.Name) + + // Readiness check + if err := flatmapHandler.WaitUntilReady(ctx); err != nil { + return fmt.Errorf("failed on map UDF readiness check, %w", err) + } + defer func() { + err = flatmapHandler.CloseConn(ctx) + if err != nil { + log.Warnw("Failed to close gRPC client conn", zap.Error(err)) + } + }() + + for index, bufferPartition := range fromBuffer { + // Populate shuffle function map + shuffleFuncMap := make(map[string]*shuffle.Shuffle) + for _, edge := range u.VertexInstance.Vertex.Spec.ToEdges { + if edge.GetToVertexPartitionCount() > 1 { + s := shuffle.NewShuffle(edge.To, edge.GetToVertexPartitionCount()) + shuffleFuncMap[fmt.Sprintf("%s:%s", edge.From, edge.To)] = s + } + } + + // create a conditional forwarder for each partition + conditionalForwarder := forwarder.GoWhere(func(keys []string, tags []string, msgId string) ([]forwarder.VertexBuffer, error) { + var result []forwarder.VertexBuffer + + // Drop message if it contains the special tag + if sharedutil.StringSliceContains(tags, dfv1.MessageTagDrop) { + return result, nil + } + + //var idString strings.Builder + // Iterate through the edges + for _, edge := range u.VertexInstance.Vertex.Spec.ToEdges { + //idString.Reset() + //idString.WriteString(edge.From) + //idString.WriteString(":") + //idString.WriteString(edge.To) + //edgeKey := idString.String() + edgeKey := fmt.Sprintf("%s:%s", edge.From, edge.To) + + // Condition to proceed for forwarding message: No conditions on edge, or message tags match edge conditions + proceed := edge.Conditions == nil || edge.Conditions.Tags == nil || len(edge.Conditions.Tags.Values) == 0 || sharedutil.CompareSlice(edge.Conditions.Tags.GetOperator(), tags, edge.Conditions.Tags.Values) + + if proceed { + // if the edge has more than one partition, shuffle the message + // else forward the message to the default partition + partitionIdx := isb.DefaultPartitionIdx + if edge.GetToVertexPartitionCount() > 1 { + if edge.ToVertexType == dfv1.VertexTypeReduceUDF { // Shuffle on keys + partitionIdx = shuffleFuncMap[edgeKey].ShuffleOnKeys(keys) + } else { // Shuffle on msgId + partitionIdx = shuffleFuncMap[edgeKey].ShuffleOnId(msgId) + } + } + + result = append(result, forwarder.VertexBuffer{ + ToVertexName: edge.To, + ToVertexPartitionIdx: partitionIdx, + }) + } + } + + return result, nil + }) + + opts := []forward.Option{} + if x := u.VertexInstance.Vertex.Spec.Limits; x != nil { + if x.ReadBatchSize != nil { + opts = append(opts, forward.WithReadBatchSize(int64(*x.ReadBatchSize))) + opts = append(opts, forward.WithUDFConcurrency(int(*x.ReadBatchSize))) + } + } + // create a forwarder for each partition + df, err := forward.NewInterStepDataForward(u.VertexInstance, readers[index], writers, conditionalForwarder, flatmapHandler, fetchWatermark, publishWatermark, idleManager, opts...) + if err != nil { + return err + } + finalWg.Add(1) + + // start the df for each partition using a go routine + go func(fromBufferPartitionName string, isdf *forward.InterStepDataForward) { + defer finalWg.Done() + log.Infow("Start processing udf messages", zap.String("isbsvc", string(u.ISBSvcType)), zap.String("from", fromBufferPartitionName), zap.Any("to", u.VertexInstance.Vertex.GetToBuffers())) + + stopped := isdf.Start() + wg := &sync.WaitGroup{} + wg.Add(1) + go func() { + defer wg.Done() + for { + <-stopped + log.Info("Forwarder stopped, exiting udf data processor for partition " + fromBufferPartitionName + "...") + return + } + }() + + <-ctx.Done() + log.Info("SIGTERM, exiting inside partition...", zap.String("partition", fromBufferPartitionName)) + isdf.Stop() + wg.Wait() + log.Info("Exited for partition...", zap.String("partition", fromBufferPartitionName)) + }(bufferPartition, df) + } + // create lag readers from buffer readers + var lagReaders []isb.LagReader + for _, reader := range readers { + lagReaders = append(lagReaders, reader) + } + + var metricsOpts []metrics.Option + metricsOpts = metrics.NewMetricsOptions(ctx, u.VertexInstance.Vertex, []metrics.HealthChecker{flatmapHandler}, lagReaders) + + ms := metrics.NewMetricsServer(u.VertexInstance.Vertex, metricsOpts...) + if shutdown, err := ms.Start(ctx); err != nil { + return fmt.Errorf("failed to start metrics server, error: %w", err) + } else { + defer func() { _ = shutdown(context.Background()) }() + } + + // wait for all the forwarders to exit + finalWg.Wait() + + // closing the publisher will only delete the keys from the store, but not the store itself + // we cannot close the store inside publisher because in some cases stores are shared between publishers + // and store itself is a separate entity that can be used by other components + for _, publisher := range publishWatermark { + err = publisher.Close() + if err != nil { + log.Errorw("Failed to close the watermark publisher", zap.Error(err)) + } + } + + // close the from vertex wm stores + // since we created the stores, we can close them + for _, wmStore := range fromVertexWmStores { + _ = wmStore.Close() + } + + // close the to vertex wm stores + // since we created the stores, we can close them + for _, wmStore := range toVertexWmStores { + _ = wmStore.Close() + } + + log.Info("All udf data processors exited...") + return nil +} diff --git a/pkg/flatmap/forward/applier/mapper.go b/pkg/flatmap/forward/applier/mapper.go new file mode 100644 index 0000000000..b3cbbfe355 --- /dev/null +++ b/pkg/flatmap/forward/applier/mapper.go @@ -0,0 +1,40 @@ +/* +Copyright 2022 The Numaproj Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package applier + +import ( + "context" + + flatmappb "github.com/numaproj/numaflow-go/pkg/apis/proto/flatmap/v1" + + "github.com/numaproj/numaflow/pkg/flatmap/types" +) + +// FlatmapApplier applies the GRPCBasedMapUDF on the stream of read messages and gives back a new message. +type FlatmapApplier interface { + // ApplyMap applies the Map UDF on a batch of N requests and streams the ResponseFlatmap response. + // It doesn't wait for the response for all the requests, before starting to send the responses back. + // It returns a channel on which any errors occurring during the processing can be propagated + ApplyMap(ctx context.Context, messageStream []*types.RequestFlatmap, writeChan chan<- *flatmappb.MapResponse) (chan struct{}, <-chan error) +} + +// ApplyFlatmapFunc utility function used to create a FlatmapApplier implementation +type ApplyFlatmapFunc func(ctx context.Context, messageStream []*types.RequestFlatmap, writeChan chan<- *flatmappb.MapResponse) (chan struct{}, <-chan error) + +func (f ApplyFlatmapFunc) ApplyMap(ctx context.Context, messageStream []*types.RequestFlatmap, writeChan chan<- *flatmappb.MapResponse) (chan struct{}, <-chan error) { + return f(ctx, messageStream, writeChan) +} diff --git a/pkg/flatmap/forward/applier/mapstreamer.go b/pkg/flatmap/forward/applier/mapstreamer.go new file mode 100644 index 0000000000..19e88862c6 --- /dev/null +++ b/pkg/flatmap/forward/applier/mapstreamer.go @@ -0,0 +1,36 @@ +/* +Copyright 2022 The Numaproj Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package applier + +import ( + "context" + + "github.com/numaproj/numaflow/pkg/isb" +) + +// MapStreamApplier applies the map stream udf on the read message and gives back response. Any UserError will be retried here, while +// InternalErr can be returned and could be retried by the callee. +type MapStreamApplier interface { + ApplyMapStream(ctx context.Context, message *isb.ReadMessage, writeMessageCh chan<- isb.WriteMessage) error +} + +// ApplyMapStreamFunc utility function used to create a MapStreamApplier implementation +type ApplyMapStreamFunc func(context.Context, *isb.ReadMessage, chan<- isb.WriteMessage) error + +func (f ApplyMapStreamFunc) ApplyMapStream(ctx context.Context, message *isb.ReadMessage, writeMessageCh chan<- isb.WriteMessage) error { + return f(ctx, message, writeMessageCh) +} diff --git a/pkg/flatmap/forward/doc.go b/pkg/flatmap/forward/doc.go new file mode 100644 index 0000000000..e2425964f0 --- /dev/null +++ b/pkg/flatmap/forward/doc.go @@ -0,0 +1,15 @@ +/* +Copyright 2022 The Numaproj Authors. +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + http://www.apache.org/licenses/LICENSE-2.0 +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +// Package forward is used for creating a data forwarder for the UDF vertex. +package forward diff --git a/pkg/flatmap/forward/forward.go b/pkg/flatmap/forward/forward.go new file mode 100644 index 0000000000..e8e1783e9a --- /dev/null +++ b/pkg/flatmap/forward/forward.go @@ -0,0 +1,832 @@ +/* +Copyright 2022 The Numaproj Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +/* +Package forward does the Read (fromBufferPartition) -> Process (map UDF) -> Forward (toBuffers) -> Ack (fromBufferPartition) loop. +*/ +package forward + +import ( + "context" + "errors" + "fmt" + "math" + "strconv" + "strings" + "sync" + "sync/atomic" + "time" + + flatmappb "github.com/numaproj/numaflow-go/pkg/apis/proto/flatmap/v1" + "go.uber.org/zap" + "k8s.io/apimachinery/pkg/util/wait" + + dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1" + "github.com/numaproj/numaflow/pkg/flatmap/forward/applier" + "github.com/numaproj/numaflow/pkg/flatmap/tracker" + "github.com/numaproj/numaflow/pkg/flatmap/types" + "github.com/numaproj/numaflow/pkg/forwarder" + "github.com/numaproj/numaflow/pkg/isb" + "github.com/numaproj/numaflow/pkg/metrics" + "github.com/numaproj/numaflow/pkg/shared/logging" + "github.com/numaproj/numaflow/pkg/watermark/fetch" + "github.com/numaproj/numaflow/pkg/watermark/publish" + "github.com/numaproj/numaflow/pkg/watermark/wmb" +) + +// InterStepDataForward forwards the data from previous step to the current step via inter-step buffer. +type InterStepDataForward struct { + // I have my reasons for overriding the default principle https://github.com/golang/go/issues/22602 + ctx context.Context + // cancelFn cancels our new context, our cancellation is little more complex and needs to be well orchestrated, hence + // we need something more than a cancel(). + cancelFn context.CancelFunc + fromBufferPartition isb.BufferReader + // toBuffers is a map of toVertex name to the toVertex's owned buffers. + toBuffers map[string][]isb.BufferWriter + FSD forwarder.ToWhichStepDecider + flatmapUDF applier.FlatmapApplier + wmFetcher fetch.Fetcher + // wmPublishers stores the vertex to publisher mapping + wmPublishers map[string]publish.Publisher + opts options + vertexName string + pipelineName string + vertexReplica int32 + // idleManager manages the idle watermark status. + idleManager wmb.IdleManager + // wmbChecker checks if the idle watermark is valid when the len(readMessage) is 0. + wmbChecker wmb.WMBChecker + responseTracker *sync.Map + requestTracker *tracker.Tracker + Shutdown +} + +// NewInterStepDataForward creates an inter-step forwarder. +func NewInterStepDataForward(vertexInstance *dfv1.VertexInstance, fromStep isb.BufferReader, toSteps map[string][]isb.BufferWriter, fsd forwarder.ToWhichStepDecider, applyUDF applier.FlatmapApplier, fetchWatermark fetch.Fetcher, publishWatermark map[string]publish.Publisher, idleManager wmb.IdleManager, opts ...Option) (*InterStepDataForward, error) { + + options := DefaultOptions() + for _, o := range opts { + if err := o(options); err != nil { + return nil, err + } + } + // creating a context here which is managed by the forwarder's lifecycle + ctx, cancel := context.WithCancel(context.Background()) + + var isdf = InterStepDataForward{ + ctx: ctx, + cancelFn: cancel, + fromBufferPartition: fromStep, + toBuffers: toSteps, + FSD: fsd, + flatmapUDF: applyUDF, + wmFetcher: fetchWatermark, + wmPublishers: publishWatermark, + // should we do a check here for the values not being null? + vertexName: vertexInstance.Vertex.Spec.Name, + pipelineName: vertexInstance.Vertex.Spec.PipelineName, + vertexReplica: vertexInstance.Replica, + idleManager: idleManager, + wmbChecker: wmb.NewWMBChecker(2), // TODO: make configurable + Shutdown: Shutdown{ + rwlock: new(sync.RWMutex), + }, + opts: *options, + responseTracker: new(sync.Map), + requestTracker: tracker.NewTracker(), + } + + // Add logger from parent ctx to child context. + isdf.ctx = logging.WithLogger(ctx, options.logger) + + if isdf.opts.enableMapUdfStream && isdf.opts.readBatchSize != 1 { + return nil, fmt.Errorf("batch size is not 1 with map UDF streaming") + } + + return &isdf, nil +} + +// Start starts reading the buffer and forwards to the next buffers. Call `Stop` to stop. +func (isdf *InterStepDataForward) Start() <-chan struct{} { + log := logging.FromContext(isdf.ctx) + stopped := make(chan struct{}) + var wg sync.WaitGroup + wg.Add(1) + go func() { + log.Info("Starting forwarder...") + // with wg approach can do more cleanup in case we need in the future. + defer wg.Done() + for { + select { + case <-isdf.ctx.Done(): + ok, err := isdf.IsShuttingDown() + if err != nil { + // ignore the error for now. + log.Errorw("Failed to check if it can shutdown", zap.Error(err)) + } + if ok { + log.Info("Shutting down...") + return + } + default: + // once context.Done() is called, we still have to try to forwardAChunk because in graceful + // shutdown the fromBufferPartition should be empty. + } + // keep doing what you are good at + isdf.forwardAChunk(isdf.ctx) + } + }() + + go func() { + wg.Wait() + // Clean up resources for buffer reader and all the writers if any. + if err := isdf.fromBufferPartition.Close(); err != nil { + log.Errorw("Failed to close buffer reader, shutdown anyways...", zap.Error(err)) + } else { + log.Infow("Closed buffer reader", zap.String("bufferFrom", isdf.fromBufferPartition.GetName())) + } + for _, buffer := range isdf.toBuffers { + for _, partition := range buffer { + if err := partition.Close(); err != nil { + log.Errorw("Failed to close partition writer, shutdown anyways...", zap.Error(err), zap.String("bufferTo", partition.GetName())) + } else { + log.Infow("Closed partition writer", zap.String("bufferTo", partition.GetName())) + } + } + } + + close(stopped) + }() + + return stopped +} + +// forwardAChunk forwards a chunk of messages from the fromBufferPartition to the toBuffers. +// It does the Read -> Process -> Forward -> Ack chain for a chunk of messages returned by the first Read call. +// It will return only if it was able to process all the message read after forwarding, barring any platform errors. +// The platform errors include buffer-full, buffer-not-reachable, etc., but does not include errors due to user code UDFs, WhereTo, etc. +// Internally, the UDF processing, write and ack are functioning asynchronously, which is done to help with a situation where +// a given message takes a long time to process, thus it does not block the other messages in this case. +// Though, we would still wait for the whole chain to complete till ack for all the messages before moving on to the +// next batch +func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { + // start time for processing forwardAChunk + start := time.Now() + + // Step 1: Read messages from the ISB + + // we read from the ISB buffer for the fromBufferPartition <= isdf.opts.readBatchSize number of messages. + // There is a chance that we have read the message and the container got forcefully terminated before processing. To provide + // at-least-once semantics for reading, during restart we will have to reprocess all unacknowledged messages. It is the + // responsibility of the Read function to do that. + readMessages, err := isdf.fromBufferPartition.Read(ctx, isdf.opts.readBatchSize) + isdf.opts.logger.Debugw("Read from buffer", zap.String("bufferFrom", isdf.fromBufferPartition.GetName()), zap.Int64("length", int64(len(readMessages)))) + if err != nil { + // TODO(stream): if we are not able to read, should we have a retry? Dont see in the code path + isdf.opts.logger.Warnw("failed to read fromBufferPartition", zap.Error(err)) + metrics.ReadMessagesError.With(map[string]string{metrics.LabelVertex: isdf.vertexName, metrics.LabelPipeline: isdf.pipelineName, metrics.LabelVertexType: string(dfv1.VertexTypeMapUDF), metrics.LabelVertexReplicaIndex: strconv.Itoa(int(isdf.vertexReplica)), metrics.LabelPartitionName: isdf.fromBufferPartition.GetName()}).Inc() + } + + // TODO(stream): check for idle watermark here Now that we are reading just a batch, we should be able + // to publish watermark similar to the current map logic. Revisit once, happy path/error path is ackDone + // process only if we have any read messages. There is a natural looping here if there is an internal error while + // reading, and we are not able to proceed. + //if len(readMessages) == 0 { + // // When the read length is zero, the write length is definitely zero too, + // // meaning there's no data to be published to the next vertex, and we consider this + // // situation as idling. + // // In order to continue propagating watermark, we will set watermark idle=true and publish it. + // // We also publish a control message if this is the first time we get this idle situation. + // // We compute the HeadIdleWMB using the given partition as the idle watermark + // var processorWMB = isdf.wmFetcher.ComputeHeadIdleWMB(isdf.fromBufferPartition.GetPartitionIdx()) + // if !isdf.wmbChecker.ValidateHeadWMB(processorWMB) { + // // validation failed, skip publishing + // isdf.opts.logger.Debugw("skip publishing idle watermark", + // zap.Int("counter", isdf.wmbChecker.GetCounter()), + // zap.Int64("offset", processorWMB.Offset), + // zap.Int64("watermark", processorWMB.Watermark), + // zap.Bool("idle", processorWMB.Idle)) + // return + // } + // + // // if the validation passed, we will publish the watermark to all the toBuffer partitions. + // for toVertexName, toVertexBuffer := range isdf.toBuffers { + // for _, partition := range toVertexBuffer { + // if p, ok := isdf.wmPublishers[toVertexName]; ok { + // idlehandler.PublishIdleWatermark(ctx, isdf.fromBufferPartition.GetPartitionIdx(), partition, p, isdf.idleManager, isdf.opts.logger, isdf.vertexName, isdf.pipelineName, dfv1.VertexTypeMapUDF, isdf.vertexReplica, wmb.Watermark(time.UnixMilli(processorWMB.Watermark))) + // } + // } + // } + // return + //} + if len(readMessages) == 0 { + return + } + logger := isdf.opts.logger + logger.Info("MYDEBUG: NO concur Bi-di ", len(readMessages)) + + // TODO(stream): see if that can be optimised by not duplicating the data slice, and passing + // We send only the dataMessages to the UDF for processing, for the non data messages, + // we just need to ack? + var dataMessages = make([]*types.RequestFlatmap, 0, len(readMessages)) + var ctrlMessageOffsets = make([]isb.Offset, 0) // for a high TPS pipeline, 0 is the most optimal value + // the readMessages itself store the offsets of the messages we read from ISB + var readOffsets = make([]isb.Offset, len(readMessages)) + for idx, m := range readMessages { + readOffsets[idx] = m.ReadOffset + if m.Kind == isb.Data { + newRequest := isdf.createNewRequest(m) + //isdf.opts.logger.Info("MYDEBUG NEW REQUEST ", newRequest.Uid) + dataMessages = append(dataMessages, newRequest) + + } else { + ctrlMessageOffsets = append(ctrlMessageOffsets, m.ReadOffset) + } + } + + // Metrics for reading data, we use ReadDataMessagesCount for calculating processing rate as well for a vertex + metrics.ReadDataMessagesCount.With(map[string]string{metrics.LabelVertex: isdf.vertexName, metrics.LabelPipeline: isdf.pipelineName, metrics.LabelVertexType: string(dfv1.VertexTypeMapUDF), metrics.LabelVertexReplicaIndex: strconv.Itoa(int(isdf.vertexReplica)), metrics.LabelPartitionName: isdf.fromBufferPartition.GetName()}).Add(float64(len(dataMessages))) + metrics.ReadMessagesCount.With(map[string]string{metrics.LabelVertex: isdf.vertexName, metrics.LabelPipeline: isdf.pipelineName, metrics.LabelVertexType: string(dfv1.VertexTypeMapUDF), metrics.LabelVertexReplicaIndex: strconv.Itoa(int(isdf.vertexReplica)), metrics.LabelPartitionName: isdf.fromBufferPartition.GetName()}).Add(float64(len(readMessages))) + + //// fetch watermark if available + //// TODO: make it async (concurrent and wait later) + //// let's track only the first element's watermark. This is important because we reassign the watermark we fetch + //// to all the elements in the batch. If we were to assign last element's watermark, we will wrongly mark on-time data as late. + //// we fetch the watermark for the partition from which we read the message. + // TODO(stream): enable to check for watermark here + //processorWM := isdf.wmFetcher.ComputeWatermark(readMessages[0].ReadOffset, isdf.fromBufferPartition.GetPartitionIdx()) + + // Step 2: UDF processing + // This involves sending the read messages to the UDF and getting the results + // To keep things asynchronous, we do not do a blocking wait on the responses from the UDF. + // For facilitating this, we use a bi-directional grpc stream connection, on which we keep sending the requests + // and then wait for the responses concurrently. + // The responses are then sent to the writer for writing to the toBuffers. + // udfRespCh is the channel on which the responses from the UDF are sent, and then this is consumed by + // the writer. + // These responses are then sent to the writer for writing to the toBuffers. + // This channel is closed when the UDF processing is ackDone, to indicate that no further processing is required. + // TODO(stream): should we keep this buffered so that on shutdown we can drain whatever is completed, + // either as ack/no ack, also to check the responsibility for close + udfRespCh := make(chan *flatmappb.MapResponse) + + // Send the input messages for processing + // The error channel returned is used to signal any errors that might have occurred during the UDF processing. + udfDoneChan, udfErrorCh := isdf.flatmapUDF.ApplyMap(ctx, dataMessages, udfRespCh) + // TODO(stream): got an error from the UDF, handle this gracefully. + go func() { + defer close(udfRespCh) + select { + // udfDoneChan indicates that we are done processing from the UDF side and do not expect any more inputs + case <-udfDoneChan: + return + + case udfErr := <-udfErrorCh: + // TODO(stream): when context is cancelled, do we want to do some different handling? + if errors.Is(udfErr, context.Canceled) || ctx.Err() != nil { + isdf.opts.logger.Infow("Context is canceled", zap.Error(udfErr)) + return + } + // We got an error while processing the UDF messages, at this point because of the nature of gRPC being + // a stream, we do not have a way to single out which request was the culprit which caused the error. + // In such a scenario to handle this we would need to replay the messages in the batch which have not been + // acked yet. So the easy way is to restart the numa container and force a reread from the ISB in the + // next cycle. + // TODO(stream): no-ack the messages which haven't been acked yet and then panic, might save + // on the replay timeout + // TODO(stream): see if there is a more graceful way in which we trigger a drain and let the + // messages which have already been processed from the UDF to complete, if that gives us a better + // performance in some way + if udfErr != nil { + isdf.opts.logger.Panic("Got an error while invoking ApplyMap", zap.Error(udfErr)) + } + } + }() + + // TODO(stream): Publish the watermark for these writers + + // Step 3: Forward to next buffer + // We keep reading on the udfRespCh to consume the messages and then forward them to the next buffer + // Further, after forwarding the messages to the buffer, the corresponding requests are streamed on the + // toAckChan so that they can be consumed for Acking to the previous buffer. + toAckChan := isdf.invokeWriter(ctx, udfRespCh) + + // Step 4: Ack the request messages to prev buffer + // We keep reading on the toAckChan to consume the messages and then ack them to the prev buffer + ackDone := isdf.invokeAck(ctx, toAckChan) + + // Wait until the Acking has been completed + <-ackDone + + //TODO(stream): is it fine if we ack the ctrlMessageOffsets in the end? + // ack the control messages, also based on some error do not ack them + if len(ctrlMessageOffsets) != 0 { + err := isdf.ackFromBuffer(ctx, ctrlMessageOffsets) + if err != nil { + return + } + } + + //TODO(stream): once processing has been completed, should we reset the tracker + // or check if any messages are left in that which have not been processed and those can be noAcked? + requestNotProcessed := isdf.requestTracker.GetItems() + // if we have requests left to process NoAck all of them + if len(requestNotProcessed) > 0 { + isdf.opts.logger.Debugw("MYDEBUG: requests left to process ", len(requestNotProcessed)) + isdf.fromBufferPartition.NoAck(ctx, readOffsets) + // reset the request tracker + isdf.requestTracker.Clear() + } + metrics.AckMessagesCount.With(map[string]string{metrics.LabelVertex: isdf.vertexName, metrics.LabelPipeline: isdf.pipelineName, metrics.LabelVertexType: string(dfv1.VertexTypeMapUDF), metrics.LabelVertexReplicaIndex: strconv.Itoa(int(isdf.vertexReplica)), metrics.LabelPartitionName: isdf.fromBufferPartition.GetName()}).Add(float64(len(dataMessages) - len(requestNotProcessed) + len(ctrlMessageOffsets))) + + // reset response tracker + isdf.responseTracker = new(sync.Map) + + isdf.opts.logger.Debugw("forwardAChunk with UDF completed", zap.Int("concurrency", isdf.opts.udfConcurrency), zap.Duration("took", time.Since(start))) +} + +// ackRoutine is a worker routine used to ack messages to the prev buffer. +// It keeps reading constantly on the ackMsgChan for any new messages, and then acks them +// Once, there are no more messages left to read on the channel, the routine exits. +func (isdf *InterStepDataForward) ackRoutine(ctx context.Context, ackMsgChan <-chan *types.RequestFlatmap, wg *sync.WaitGroup) { + defer wg.Done() + for response := range ackMsgChan { + //isdf.opts.logger.Info("MYDEBUG: GOT TO ACK ", response.Uid) + if err := isdf.ackFromBufferSingle(ctx, response.Request.ReadOffset); err != nil { + isdf.opts.logger.Error("MYDEBUG: ERROR IN ACK ", zap.Error(err)) + // TODO(stream): we have retried in the ackFromBuffer, should we trigger drain here then? + isdf.requestTracker.RemoveRequest(response.Uid) + metrics.AckMessageError.With(map[string]string{metrics.LabelVertex: isdf.vertexName, metrics.LabelPipeline: isdf.pipelineName, metrics.LabelVertexType: string(dfv1.VertexTypeMapUDF), metrics.LabelVertexReplicaIndex: strconv.Itoa(int(isdf.vertexReplica)), metrics.LabelPartitionName: isdf.fromBufferPartition.GetName()}).Inc() + return + } + isdf.requestTracker.RemoveRequest(response.Uid) + //isdf.opts.logger.Info("MYDEBUG: DONE ACK ", response.Uid) + } +} + +// invokeAck is the function used to ack the requests to the prev buffer. +// it orchestrates a ack routine pool, consumes the responses from the writer and then sends it to the pool for +// acking. Each worker in the pool keeps reading on the input channel for any new message, and then acks it. +// Once all the pool workers exit, we consider the acking jobs to be completed and then close the done to indicate +// this. +// TODO(stream): check if the ack path can be optimised as now we are writing one message per worker, instead of sending a batch for writing. +func (isdf *InterStepDataForward) invokeAck(ctx context.Context, ackMsgChan <-chan *types.RequestFlatmap) (doneChan chan struct{}) { + doneChan = make(chan struct{}) + go func() { + defer close(doneChan) + group := sync.WaitGroup{} + for i := 0; i < int(isdf.opts.readBatchSize); i++ { + group.Add(1) + go isdf.ackRoutine(ctx, ackMsgChan, &group) + } + group.Wait() + }() + return doneChan +} + +// writeRoutine is a worker routine used to forward messages to the next buffer. +// It keeps reading constantly on the udfRespCh for any new messages, and then forwards it to the correct +// next buffer according to the conditional logic. +// Once, there are no more messages left to read on the channel, the routine exits. +// If there is an error in the UDF processing the udfRespCh is closed, so the workers should exit +func (isdf *InterStepDataForward) writeRoutine(ctx context.Context, udfRespCh <-chan *flatmappb.MapResponse, ackChan chan<- *types.RequestFlatmap, wg *sync.WaitGroup) { + addCount := 0 + defer func(int) { + metrics.UDFWriteMessagesCount.With(map[string]string{metrics.LabelVertex: isdf.vertexName, metrics.LabelPipeline: isdf.pipelineName, metrics.LabelVertexType: string(dfv1.VertexTypeMapUDF), metrics.LabelVertexReplicaIndex: strconv.Itoa(int(isdf.vertexReplica)), metrics.LabelPartitionName: isdf.fromBufferPartition.GetName()}).Add(float64(addCount)) + wg.Done() + }(addCount) + for response := range udfRespCh { + uid := response.Result.GetUuid() + trackedRequest, ok := isdf.requestTracker.GetRequest(uid) + // TODO(stream): check what should be path for !ok, which means that we got a UUID + // which has already been deleted from the tracker/ or never added in the first place + // can this even happen though if messages are ordered and we only have a single routine processing it? + if !ok { + isdf.opts.logger.Error("MYDEBUG: MESSAGE NOT IN TRACKER ", response.Result.GetTotal(), uid) + } + parsedResp, requestDone, total := isdf.parseMapResponse(response, trackedRequest) + // if the response has the AckIt field = true, that indicates the end of the processing for + // a given message. In this case send the parent request for Acking and continue + // if requestDone (EOR == true) and total == 0 it means there were no responses + // expected for this request, hence we can directly ackIt, no need to write + if requestDone && total == 0 { + ackChan <- trackedRequest + continue + } + // If AckIt is not set, it is a data response, hence forward it to the next buffer + var messageToStep = make(map[string][]isb.Message) + for toVertex := range isdf.toBuffers { + // over allocating to have a predictable pattern + messageToStep[toVertex] = make([]isb.Message, len(isdf.toBuffers[toVertex])) + } + writeMessage := parsedResp + if err := isdf.forwardToBuffers(ctx, writeMessage, trackedRequest.Request, messageToStep); err != nil { + // As we have re-tried already to forward to the buffer, we should not be trying it again. + // TODO(stream): so we should just exit and send this for no-ack? or exit + isdf.opts.logger.Error("MYDEBUG: NEW ERROR IN WRITE", zap.Error(err)) + } + addCount++ + // Update the response counter for the given UUID (request) + idxNum := isdf.updateCounter(uid) + // If the counter has reached the total number of responses expected, we can safely + // send the parent request for Acking + if idxNum == int64(total) { + ackChan <- trackedRequest + } + } +} + +func (isdf *InterStepDataForward) updateCounter(key string) int64 { + val, _ := isdf.responseTracker.LoadOrStore(key, new(int64)) + ptr := val.(*int64) + newVal := atomic.AddInt64(ptr, 1) + return newVal +} + +// invokeWriter is the function to forward the responses to the next buffer. +// it orchestrates a writer routine pool, consumes the responses from the UDF and then sends it to the pool for +// forwarding. Each worker in the pool keeps reading on the input channel for any new message, and then forwards it. +// After forwarding the message we send the corresponding parent request for Acking via the ackChan +// Once all the pool workers exit, we consider the writing jobs to be completed and then close the ackChan to indicate +// this further. +// TODO(stream): check if the write path can be optimised as now we are writing one message per worker, instead of sending a batch for writing. +func (isdf *InterStepDataForward) invokeWriter(ctx context.Context, writeMessageCh <-chan *flatmappb.MapResponse) <-chan *types.RequestFlatmap { + // ackChan is used to stream the ReadMessage which need to be Acked + // TODO(stream): if we want to send something for noAck explicitly, might want to use types.AckMsgFlatmap + ackChan := make(chan *types.RequestFlatmap) + go func() { + // close to indicate that no further messages left to ack + defer close(ackChan) + group := sync.WaitGroup{} + // start the pool + for i := 0; i < int(isdf.opts.readBatchSize); i++ { + group.Add(1) + go isdf.writeRoutine(ctx, writeMessageCh, ackChan, &group) + } + // wait until all the writeRoutines are done + group.Wait() + }() + return ackChan +} + +func (isdf *InterStepDataForward) forwardToBuffers(ctx context.Context, writeMessages *isb.WriteMessage, readMessage *isb.ReadMessage, messageToStep map[string][]isb.Message) error { + if writeMessages == nil { + return nil + } + if err := isdf.whereToStep(writeMessages, messageToStep, readMessage); err != nil { + isdf.opts.logger.Errorw("failed in whereToStep", zap.Error(err)) + return err + } + // forward the messages to the edge buffer (could be multiple edges) + _, err := isdf.writeToBuffers(ctx, messageToStep) + if err != nil { + isdf.opts.logger.Errorw("failed to write to toBuffers", zap.Error(err)) + return err + } + return nil +} + +// ackFromBufferSingle acknowledges an array of offsets back to fromBufferPartition and is a blocking call or until shutdown has been initiated. +func (isdf *InterStepDataForward) ackFromBufferSingle(ctx context.Context, offset isb.Offset) error { + //for { + // select { + // case <-ctx.Done(): + // return ctx.Err() + // default: + // if ok, _ := isdf.IsShuttingDown(); ok { + // ackErr := fmt.Errorf("AckFromBuffer, Stop called while stuck on an internal error") + // return ackErr + // } + // err := offset.AckIt() + // if err != nil { + // // TODO(stream): add metrics here + // isdf.opts.logger.Errorw("Failed to ack message", zap.Error(err)) + // // If the error is related to nats/jetstream, we skip it because it might end up with infinite ack retries. + // // Skipping those errors to let the whole read/write/ack for loop to restart from reading, to pick up those + // // redelivered messages. + // if !strings.HasPrefix(err.Error(), "nats:") { + // return err + // } + // continue + // } else { + // return nil + // } + // } + //} + + var ackRetryBackOff = wait.Backoff{ + Factor: 1, + Jitter: 0.1, + Steps: math.MaxInt, + Duration: time.Millisecond * 10, + } + //var ackOffsets = []isb.Offset{offset} + attempt := 0 + ctxClosedErr := wait.ExponentialBackoff(ackRetryBackOff, func() (done bool, err error) { + //errs := isdf.fromBufferPartition.Ack(ctx, ackOffsets) + errs := offset.AckIt() + attempt += 1 + //if errs != nil && errs[0] != nil { + if errs != nil { + // TODO(stream): add metrics here + //var summarizedErr = errs[0].Error() + var summarizedErr = errs.Error() + isdf.opts.logger.Errorw("Failed to ack from buffer, retrying", zap.Any("errors", summarizedErr), zap.Int("attempt", attempt)) + // no point retrying if ctx.Done has been invoked + select { + case <-ctx.Done(): + // no point in retrying after we have been asked to stop. + return false, ctx.Err() + default: + if ok, _ := isdf.IsShuttingDown(); ok { + ackErr := fmt.Errorf("AckFromBuffer, Stop called while stuck on an internal error, %v", summarizedErr) + return false, ackErr + } + return false, nil + } + } else { + return true, nil + } + }) + + if ctxClosedErr != nil { + isdf.opts.logger.Errorw("Context closed while waiting to ack messages inside forward", zap.Error(ctxClosedErr)) + } + return ctxClosedErr +} + +func (isdf *InterStepDataForward) ackFromBuffer(ctx context.Context, offsets []isb.Offset) error { + var ackRetryBackOff = wait.Backoff{ + Factor: 1, + Jitter: 0.1, + Steps: math.MaxInt, + Duration: time.Millisecond * 10, + } + var ackOffsets = offsets + attempt := 0 + + ctxClosedErr := wait.ExponentialBackoff(ackRetryBackOff, func() (done bool, err error) { + errs := isdf.fromBufferPartition.Ack(ctx, ackOffsets) + attempt += 1 + summarizedErr := errorArrayToMap(errs) + var failedOffsets []isb.Offset + if len(summarizedErr) > 0 { + isdf.opts.logger.Errorw("Failed to ack from buffer, retrying", zap.Any("errors", summarizedErr), zap.Int("attempt", attempt)) + // no point retrying if ctx.Done has been invoked + select { + case <-ctx.Done(): + // no point in retrying after we have been asked to stop. + return false, ctx.Err() + default: + // retry only the failed offsets + for i, offset := range ackOffsets { + if errs[i] != nil { + failedOffsets = append(failedOffsets, offset) + } + } + ackOffsets = failedOffsets + if ok, _ := isdf.IsShuttingDown(); ok { + ackErr := fmt.Errorf("AckFromBuffer, Stop called while stuck on an internal error, %v", summarizedErr) + return false, ackErr + } + return false, nil + } + } else { + return true, nil + } + }) + + if ctxClosedErr != nil { + isdf.opts.logger.Errorw("Context closed while waiting to ack messages inside forward", zap.Error(ctxClosedErr)) + } + + return ctxClosedErr +} + +// writeToBuffers is a blocking call until all the messages have be forwarded to all the toBuffers, or a shutdown +// has been initiated while we are stuck looping on an InternalError. +func (isdf *InterStepDataForward) writeToBuffers( + ctx context.Context, messageToStep map[string][]isb.Message, +) (writeOffsets map[string][][]isb.Offset, err error) { + // messageToStep contains all the to buffers, so the messages could be empty (conditional forwarding). + // So writeOffsets also contains all the to buffers, but the returned offsets might be empty. + + //writeOffsets = make(map[string][][]isb.Offset) + //for toVertexName, toVertexMessages := range messageToStep { + // writeOffsets[toVertexName] = make([][]isb.Offset, len(toVertexMessages)) + //} + + for toVertexName, toVertexBuffer := range isdf.toBuffers { + for index, partition := range toVertexBuffer { + _, err = isdf.writeToBuffer(ctx, partition, messageToStep[toVertexName][index]) + if err != nil { + return nil, err + } + } + } + return nil, nil +} + +// writeToBuffer forwards an array of messages to a single buffer and is a blocking call or until shutdown has been initiated. +func (isdf *InterStepDataForward) writeToBuffer(ctx context.Context, toBufferPartition isb.BufferWriter, msg isb.Message) (writeOffsets isb.Offset, err error) { + var ( + //totalCount int + writeCount int + writeBytes float64 + ) + //totalCount = len(messages) + //writeOffsets = make([]isb.Offset, 0, totalCount) + + for { + // EXTRA + //var _writeOffsets []isb.Offset = nil + //var errs []error = nil + _writeOffsets, errs := toBufferPartition.WriteNew(ctx, msg) + // Note: this is an unwanted memory allocation during a happy path. We want only minimal allocation since using failedMessages is an unlikely path. + var failedMessages isb.Message + needRetry := false + //for idx, msg := range messages { + // EXTRA + //if err != nil { + if err != nil { + // ATM there are no user-defined errors during write, all are InternalErrors. + // Non retryable error, drop the message. Non retryable errors are only returned + // when the buffer is full and the user has set the buffer full strategy to + // DiscardLatest or when the message is duplicate. + if errors.As(err, &isb.NonRetryableBufferWriteErr{}) { + metrics.DropMessagesCount.With(map[string]string{ + metrics.LabelVertex: isdf.vertexName, + metrics.LabelPipeline: isdf.pipelineName, + metrics.LabelVertexType: string(dfv1.VertexTypeSink), + metrics.LabelVertexReplicaIndex: strconv.Itoa(int(isdf.vertexReplica)), + metrics.LabelPartitionName: toBufferPartition.GetName(), + metrics.LabelReason: err.Error(), + }).Inc() + + metrics.DropBytesCount.With(map[string]string{ + metrics.LabelVertex: isdf.vertexName, + metrics.LabelPipeline: isdf.pipelineName, + metrics.LabelVertexType: string(dfv1.VertexTypeSink), + metrics.LabelVertexReplicaIndex: strconv.Itoa(int(isdf.vertexReplica)), + metrics.LabelPartitionName: toBufferPartition.GetName(), + metrics.LabelReason: err.Error(), + }).Add(float64(len(msg.Payload))) + + isdf.opts.logger.Infow("Dropped message", zap.String("reason", err.Error()), zap.String("partition", toBufferPartition.GetName()), zap.String("vertex", isdf.vertexName), zap.String("pipeline", isdf.pipelineName)) + } else { + needRetry = true + // we retry only failed messages + failedMessages = msg + metrics.WriteMessagesError.With(map[string]string{metrics.LabelVertex: isdf.vertexName, metrics.LabelPipeline: isdf.pipelineName, metrics.LabelVertexType: string(dfv1.VertexTypeMapUDF), metrics.LabelVertexReplicaIndex: strconv.Itoa(int(isdf.vertexReplica)), metrics.LabelPartitionName: toBufferPartition.GetName()}).Inc() + // a shutdown can break the blocking loop caused due to InternalErr + if ok, _ := isdf.IsShuttingDown(); ok { + metrics.PlatformError.With(map[string]string{metrics.LabelVertex: isdf.vertexName, metrics.LabelPipeline: isdf.pipelineName, metrics.LabelVertexType: string(dfv1.VertexTypeMapUDF), metrics.LabelVertexReplicaIndex: strconv.Itoa(int(isdf.vertexReplica))}).Inc() + return writeOffsets, fmt.Errorf("writeToBuffer failed, Stop called while stuck on an internal error with failed messages: %v", errs) + } + } + } else { + writeCount++ + writeBytes += float64(len(msg.Payload)) + // we support write offsets only for jetstream + if _writeOffsets != nil { + writeOffsets = _writeOffsets + } + } + //} + + if needRetry { + isdf.opts.logger.Errorw("Retrying failed messages", + zap.Any("errors", errs.Error()), + zap.String(metrics.LabelPipeline, isdf.pipelineName), + zap.String(metrics.LabelVertex, isdf.vertexName), + zap.String(metrics.LabelPartitionName, toBufferPartition.GetName()), + ) + // set messages to failed for the retry + msg = failedMessages + // TODO: implement retry with backoff etc. + time.Sleep(isdf.opts.retryInterval) + } else { + break + } + } + + //metrics.WriteMessagesCount.With(map[string]string{metrics.LabelVertex: isdf.vertexName, metrics.LabelPipeline: isdf.pipelineName, metrics.LabelVertexType: string(dfv1.VertexTypeMapUDF), metrics.LabelVertexReplicaIndex: strconv.Itoa(int(isdf.vertexReplica)), metrics.LabelPartitionName: toBufferPartition.GetName()}).Add(float64(writeCount)) + //metrics.WriteBytesCount.With(map[string]string{metrics.LabelVertex: isdf.vertexName, metrics.LabelPipeline: isdf.pipelineName, metrics.LabelVertexType: string(dfv1.VertexTypeMapUDF), metrics.LabelVertexReplicaIndex: strconv.Itoa(int(isdf.vertexReplica)), metrics.LabelPartitionName: toBufferPartition.GetName()}).Add(writeBytes) + return writeOffsets, nil +} + +// whereToStep executes the WhereTo interfaces and then updates the to step's writeToBuffers buffer. +func (isdf *InterStepDataForward) whereToStep(writeMessage *isb.WriteMessage, messageToStep map[string][]isb.Message, readMessage *isb.ReadMessage) error { + // call WhereTo and drop it on errors + to, err := isdf.FSD.WhereTo(writeMessage.Keys, writeMessage.Tags, writeMessage.ID) + if err != nil { + isdf.opts.logger.Errorw("failed in whereToStep", zap.Error(isb.MessageWriteErr{Name: isdf.fromBufferPartition.GetName(), Header: readMessage.Header, Body: readMessage.Body, Message: fmt.Sprintf("WhereTo failed, %s", err)})) + // a shutdown can break the blocking loop caused due to InternalErr + if ok, _ := isdf.IsShuttingDown(); ok { + err := fmt.Errorf("whereToStep, Stop called while stuck on an internal error, %v", err) + metrics.PlatformError.With(map[string]string{metrics.LabelVertex: isdf.vertexName, metrics.LabelPipeline: isdf.pipelineName, metrics.LabelVertexType: string(dfv1.VertexTypeMapUDF), metrics.LabelVertexReplicaIndex: strconv.Itoa(int(isdf.vertexReplica))}).Inc() + return err + } + return err + } + + for _, t := range to { + if _, ok := messageToStep[t.ToVertexName]; !ok { + isdf.opts.logger.Errorw("failed in whereToStep", zap.Error(isb.MessageWriteErr{Name: isdf.fromBufferPartition.GetName(), Header: readMessage.Header, Body: readMessage.Body, Message: fmt.Sprintf("no such destination (%s)", t.ToVertexName)})) + } + messageToStep[t.ToVertexName][t.ToVertexPartitionIdx] = writeMessage.Message + } + return nil +} + +// errorArrayToMap summarizes an error array to map +func errorArrayToMap(errs []error) map[string]int64 { + result := make(map[string]int64) + for _, err := range errs { + if err != nil { + result[err.Error()]++ + } + } + return result +} + +func (isdf *InterStepDataForward) createNewRequest(msg *isb.ReadMessage) *types.RequestFlatmap { + // Add the request to the tracker, and get the unique UUID corresponding to it + return isdf.requestTracker.AddRequest(msg) +} + +// parseMapResponse takes a proto response from the gRPC and converts this into a ResponseFlatmap type, +// this also checks if this was a special EOR response, in such a case we indicate that the request corresponding +// to the response can be safely removed from the tracker. +func (isdf *InterStepDataForward) parseMapResponse(resp *flatmappb.MapResponse, trackedRequest *types.RequestFlatmap) (*isb.WriteMessage, bool, int32) { + result := resp.Result + eor := result.GetEOR() + total := result.GetTotal() + + parentRequest := trackedRequest.Request + // Request has completed remove from the tracker module + if eor == true { + return nil, true, total + //return &types.ResponseFlatmap{ + // //ParentMessage: parentRequest, + // Uid: uid, + // RespMessage: nil, + // AckIt: true, + // Total: int64(resp.Result.Total), + //}, true, uid + } + keys := result.GetKeys() + taggedMessage := &isb.WriteMessage{ + Message: isb.Message{ + Header: isb.Header{ + MessageInfo: parentRequest.MessageInfo, + // We need this to be unique so that the ISB can execute its Dedup logic + // this ID should be such that even when the same response is processed and received + // again from the UDF, we still assign it the same ID. + // The ID here will be a concat of the three values + // parentRequest.ReadOffset - vertexName - result.Index + // + // ReadOffset - Will be the read offset of the request which corresponds to this response. + // We have this stored in our tracker. + // + // VertexName - the name of the vertex from which this response is generated, this is + // important to ensure that we can differentiate between messages emitted from 2 map vertices + // + // Result Index - This parameter is added on the SDK side. + // We add the index of the message from the messages slice to the individual response. + // TODO(stream): explore if there can be more robust ways to do this + ID: getMessageId(trackedRequest.ReadOffset.String(), isdf.vertexName, result.GetIndex()), + Keys: keys, + }, + Body: isb.Body{ + Payload: result.GetValue(), + }, + }, + Tags: result.GetTags(), + } + return taggedMessage, false, total +} + +func getMessageId(offset string, vertexName string, index string) string { + var idString strings.Builder + idString.WriteString(offset) + idString.WriteString("-") + idString.WriteString(vertexName) + idString.WriteString("-") + idString.WriteString(index) + return idString.String() +} diff --git a/pkg/flatmap/forward/forward_test.go b/pkg/flatmap/forward/forward_test.go new file mode 100644 index 0000000000..17fd2415a1 --- /dev/null +++ b/pkg/flatmap/forward/forward_test.go @@ -0,0 +1,1741 @@ +/* +Copyright 2022 The Numaproj Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package forward + +import ( + "context" + "fmt" + "strings" + "sync" + "testing" + "time" + + "github.com/prometheus/client_golang/prometheus/testutil" + "github.com/stretchr/testify/assert" + "go.uber.org/goleak" + + dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1" + "github.com/numaproj/numaflow/pkg/forwarder" + "github.com/numaproj/numaflow/pkg/isb" + "github.com/numaproj/numaflow/pkg/isb/stores/simplebuffer" + "github.com/numaproj/numaflow/pkg/isb/testutils" + "github.com/numaproj/numaflow/pkg/metrics" + "github.com/numaproj/numaflow/pkg/shared/kvs" + "github.com/numaproj/numaflow/pkg/shared/logging" + udfapplier "github.com/numaproj/numaflow/pkg/udf/rpc" + "github.com/numaproj/numaflow/pkg/watermark/entity" + "github.com/numaproj/numaflow/pkg/watermark/generic" + "github.com/numaproj/numaflow/pkg/watermark/publish" + wmstore "github.com/numaproj/numaflow/pkg/watermark/store" + "github.com/numaproj/numaflow/pkg/watermark/wmb" +) + +const ( + testPipelineName = "testPipeline" + testProcessorEntity = "publisherTestPod" + publisherKeyspace = testPipelineName + "_" + testProcessorEntity + "_%s" +) + +var ( + testStartTime = time.Unix(1636470000, 0).UTC() + testSourceWatermark = time.Unix(1636460000, 0).UTC() + testWMBWatermark = time.Unix(1636470000, 0).UTC() +) + +type testForwardFetcher struct { + // for forward_test.go only +} + +func TestMain(m *testing.M) { + goleak.VerifyTestMain(m) +} + +func (t *testForwardFetcher) ComputeWatermark(offset isb.Offset, partition int32) wmb.Watermark { + return t.getWatermark() +} + +// getWatermark uses current time as the watermark because we want to make sure +// the test publisher is publishing watermark +func (t *testForwardFetcher) getWatermark() wmb.Watermark { + return wmb.Watermark(testSourceWatermark) +} + +func (t *testForwardFetcher) ComputeHeadIdleWMB(int32) wmb.WMB { + // won't be used + return wmb.WMB{} +} + +type myForwardTest struct { +} + +func (f myForwardTest) WhereTo(_ []string, _ []string, s string) ([]forwarder.VertexBuffer, error) { + return []forwarder.VertexBuffer{{ + ToVertexName: "to1", + ToVertexPartitionIdx: 0, + }}, nil +} + +func (f myForwardTest) ApplyMap(ctx context.Context, message *isb.ReadMessage) ([]*isb.WriteMessage, error) { + return testutils.CopyUDFTestApply(ctx, message) +} + +func (f myForwardTest) ApplyMapStream(ctx context.Context, message *isb.ReadMessage, writeMessageCh chan<- isb.WriteMessage) error { + return testutils.CopyUDFTestApplyStream(ctx, message, writeMessageCh) +} + +func TestNewInterStepDataForward(t *testing.T) { + tests := []struct { + name string + batchSize int64 + streamEnabled bool + }{ + { + name: "stream_forward", + batchSize: 1, + streamEnabled: true, + }, + { + name: "batch_forward", + batchSize: 10, + streamEnabled: false, + }, + } + for _, tt := range tests { + t.Run(tt.name+"_basic", func(t *testing.T) { + metricsReset() + batchSize := tt.batchSize + fromStep := simplebuffer.NewInMemoryBuffer("from", 5*batchSize, 0) + to11 := simplebuffer.NewInMemoryBuffer("to1-1", 2*batchSize, 0) + to12 := simplebuffer.NewInMemoryBuffer("to1-2", 2*batchSize, 1) + toSteps := map[string][]isb.BufferWriter{ + "to1": {to11, to12}, + } + + vertex := &dfv1.Vertex{Spec: dfv1.VertexSpec{ + PipelineName: "testPipeline", + AbstractVertex: dfv1.AbstractVertex{ + Name: "testVertex", + }, + }} + vertexInstance := &dfv1.VertexInstance{ + Vertex: vertex, + Replica: 0, + } + + ctx, cancel := context.WithTimeout(context.Background(), time.Second*10) + defer cancel() + + writeMessages := testutils.BuildTestWriteMessages(4*batchSize, testStartTime, nil) + + fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps) + + idleManager, _ := wmb.NewIdleManager(1, len(toSteps)) + f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, &mySourceForwardTestRoundRobin{}, myForwardTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled)) + + assert.NoError(t, err) + assert.False(t, to11.IsFull()) + assert.False(t, to12.IsFull()) + + assert.True(t, to11.IsEmpty()) + assert.True(t, to12.IsEmpty()) + + stopped := f.Start() + // write some data + _, errs := fromStep.Write(ctx, writeMessages[0:batchSize]) + assert.Equal(t, make([]error, batchSize), errs) + + var updatedBatchSize int + if tt.batchSize > 1 { + updatedBatchSize = int(batchSize / 2) + } else { + updatedBatchSize = int(batchSize) + } + // read some data + readMessages, err := to11.Read(ctx, int64(updatedBatchSize)) + assert.NoError(t, err, "expected no error") + assert.Len(t, readMessages, updatedBatchSize) + for i, j := 0, 0; i < updatedBatchSize; i, j = i+1, j+2 { + assert.Equal(t, []interface{}{writeMessages[j].MessageInfo}, []interface{}{readMessages[i].MessageInfo}) + assert.Equal(t, []interface{}{writeMessages[j].Kind}, []interface{}{readMessages[i].Kind}) + assert.Equal(t, []interface{}{writeMessages[j].Keys}, []interface{}{readMessages[i].Keys}) + assert.Equal(t, []interface{}{writeMessages[j].Body}, []interface{}{readMessages[i].Body}) + } + + if tt.batchSize > 1 { + readMessages, err = to12.Read(ctx, int64(updatedBatchSize)) + assert.NoError(t, err, "expected no error") + assert.Len(t, readMessages, updatedBatchSize) + for i, j := 0, 1; i < updatedBatchSize; i, j = i+1, j+2 { + assert.Equal(t, []interface{}{writeMessages[j].MessageInfo}, []interface{}{readMessages[i].MessageInfo}) + assert.Equal(t, []interface{}{writeMessages[j].Kind}, []interface{}{readMessages[i].Kind}) + assert.Equal(t, []interface{}{writeMessages[j].Keys}, []interface{}{readMessages[i].Keys}) + assert.Equal(t, []interface{}{writeMessages[j].Body}, []interface{}{readMessages[i].Body}) + } + } + validateMetrics(t, batchSize) + // write some data + _, errs = fromStep.Write(ctx, writeMessages[batchSize:4*batchSize]) + assert.Equal(t, make([]error, 3*batchSize), errs) + + f.Stop() + time.Sleep(1 * time.Millisecond) + // only for shutdown will work as from buffer is not empty + f.ForceStop() + + <-stopped + + }) + // Explicitly tests the case where we forward to all buffers + t.Run(tt.name+"_toAll", func(t *testing.T) { + batchSize := tt.batchSize + fromStep := simplebuffer.NewInMemoryBuffer("from", 10*batchSize, 0) + to11 := simplebuffer.NewInMemoryBuffer("to1-1", 2*batchSize, 0) + to12 := simplebuffer.NewInMemoryBuffer("to1-2", 2*batchSize, 1) + + to21 := simplebuffer.NewInMemoryBuffer("to2-1", 2*batchSize, 0) + to22 := simplebuffer.NewInMemoryBuffer("to2-2", 2*batchSize, 1) + + toSteps := map[string][]isb.BufferWriter{ + "to1": {to11, to12}, + "to2": {to21, to22}, + } + ctx, cancel := context.WithTimeout(context.Background(), time.Second*10) + defer cancel() + + writeMessages := testutils.BuildTestWriteMessages(4*batchSize, testStartTime, nil) + + vertex := &dfv1.Vertex{Spec: dfv1.VertexSpec{ + PipelineName: "testPipeline", + AbstractVertex: dfv1.AbstractVertex{ + Name: "testVertex", + }, + }} + + vertexInstance := &dfv1.VertexInstance{ + Vertex: vertex, + Replica: 0, + } + + fetchWatermark := &testForwardFetcher{} + toVertexWmStores := buildWatermarkStores(toSteps) + publishWatermark, otStores := buildPublisherMapAndOTStoreFromWmStores(toSteps, toVertexWmStores) + + defer func() { + for _, p := range publishWatermark { + _ = p.Close() + } + }() + + defer func() { + for _, store := range toVertexWmStores { + _ = store.Close() + } + }() + + idleManager, _ := wmb.NewIdleManager(1, len(toSteps)) + f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, &myForwardToAllTest{}, &myForwardToAllTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled)) + + assert.NoError(t, err) + assert.False(t, to11.IsFull()) + assert.False(t, to12.IsFull()) + + assert.True(t, to11.IsEmpty()) + assert.True(t, to12.IsEmpty()) + + stopped := f.Start() + // write some data + _, errs := fromStep.Write(ctx, writeMessages[0:batchSize]) + assert.Equal(t, make([]error, batchSize), errs) + + var updatedBatchSize int + if tt.batchSize > 1 { + updatedBatchSize = int(batchSize / 2) + } else { + updatedBatchSize = int(batchSize) + } + // read some data + readMessages, err := to11.Read(ctx, int64(updatedBatchSize)) + assert.NoError(t, err, "expected no error") + assert.Len(t, readMessages, updatedBatchSize) + for i, j := 0, 0; i < updatedBatchSize; i, j = i+1, j+2 { + assert.Equal(t, []interface{}{writeMessages[j].MessageInfo}, []interface{}{readMessages[i].MessageInfo}) + assert.Equal(t, []interface{}{writeMessages[j].Kind}, []interface{}{readMessages[i].Kind}) + assert.Equal(t, []interface{}{writeMessages[j].Keys}, []interface{}{readMessages[i].Keys}) + assert.Equal(t, []interface{}{writeMessages[j].Body}, []interface{}{readMessages[i].Body}) + } + + if tt.batchSize > 1 { + // read some data + readMessages, err = to12.Read(ctx, int64(updatedBatchSize)) + assert.NoError(t, err, "expected no error") + assert.Len(t, readMessages, updatedBatchSize) + for i, j := 0, 1; i < updatedBatchSize; i, j = i+1, j+2 { + assert.Equal(t, []interface{}{writeMessages[j].MessageInfo}, []interface{}{readMessages[i].MessageInfo}) + assert.Equal(t, []interface{}{writeMessages[j].Kind}, []interface{}{readMessages[i].Kind}) + assert.Equal(t, []interface{}{writeMessages[j].Keys}, []interface{}{readMessages[i].Keys}) + assert.Equal(t, []interface{}{writeMessages[j].Body}, []interface{}{readMessages[i].Body}) + } + } + // write some data + _, errs = fromStep.Write(ctx, writeMessages[batchSize:4*batchSize]) + assert.Equal(t, make([]error, 3*batchSize), errs) + + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + otKeys1, _ := otStores["to1"].GetAllKeys(ctx) + loop: + for { + select { + case <-ctx.Done(): + assert.Fail(t, "context cancelled while waiting to get a valid watermark") + break loop + default: + if len(otKeys1) == 0 { + otKeys1, _ = otStores["to1"].GetAllKeys(ctx) + time.Sleep(time.Millisecond * 10) + } else { + // NOTE: in this test we only have one processor to publish + // so len(otKeys) should always be 1 + otKeys1, _ = otStores["to1"].GetAllKeys(ctx) + otValue1, _ := otStores["to1"].GetValue(ctx, otKeys1[0]) + otDecode1, _ := wmb.DecodeToWMB(otValue1) + if tt.batchSize > 1 && !otDecode1.Idle { + break loop + } else if tt.batchSize == 1 && otDecode1.Idle { + break loop + } else { + time.Sleep(time.Millisecond * 10) + } + } + } + + } + }() + + wg.Add(1) + go func() { + defer wg.Done() + otKeys2, _ := otStores["to2"].GetAllKeys(ctx) + loop: + for { + select { + case <-ctx.Done(): + assert.Fail(t, "context cancelled while waiting to get a valid watermark") + break loop + default: + if len(otKeys2) == 0 { + otKeys2, _ = otStores["to2"].GetAllKeys(ctx) + time.Sleep(time.Millisecond * 10) + } else { + // NOTE: in this test we only have one processor to publish + // so len(otKeys) should always be 1 + otKeys2, _ = otStores["to2"].GetAllKeys(ctx) + otValue2, _ := otStores["to2"].GetValue(ctx, otKeys2[0]) + otDecode1, _ := wmb.DecodeToWMB(otValue2) + // if the batch size is 1, then the watermark should be idle because + // one of partitions will be idle + if tt.batchSize > 1 && !otDecode1.Idle { + break loop + } else if tt.batchSize == 1 && otDecode1.Idle { + break loop + } else { + time.Sleep(time.Millisecond * 10) + } + } + } + + } + }() + wg.Wait() + f.Stop() + <-stopped + }) + // Explicitly tests the case where we drop all events + t.Run(tt.name+"_dropAll", func(t *testing.T) { + batchSize := tt.batchSize + fromStep := simplebuffer.NewInMemoryBuffer("from", 5*batchSize, 0) + to11 := simplebuffer.NewInMemoryBuffer("to1-1", 2*batchSize, 0) + to12 := simplebuffer.NewInMemoryBuffer("to1-2", 2*batchSize, 1) + + to21 := simplebuffer.NewInMemoryBuffer("to2-1", 2*batchSize, 0) + to22 := simplebuffer.NewInMemoryBuffer("to2-2", 2*batchSize, 1) + toSteps := map[string][]isb.BufferWriter{ + "to1": {to11, to12}, + "to2": {to21, to22}, + } + ctx, cancel := context.WithTimeout(context.Background(), time.Second*10) + defer cancel() + + writeMessages := testutils.BuildTestWriteMessages(4*batchSize, testStartTime, nil) + + vertex := &dfv1.Vertex{Spec: dfv1.VertexSpec{ + PipelineName: "testPipeline", + AbstractVertex: dfv1.AbstractVertex{ + Name: "testVertex", + }, + }} + + vertexInstance := &dfv1.VertexInstance{ + Vertex: vertex, + Replica: 0, + } + + fetchWatermark := &testForwardFetcher{} + toVertexWmStores := buildWatermarkStores(toSteps) + publishWatermark, otStores := buildPublisherMapAndOTStoreFromWmStores(toSteps, toVertexWmStores) + + defer func() { + for _, p := range publishWatermark { + _ = p.Close() + } + }() + + defer func() { + for _, store := range toVertexWmStores { + _ = store.Close() + } + }() + + idleManager, _ := wmb.NewIdleManager(1, len(toSteps)) + f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardDropTest{}, myForwardDropTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled)) + + assert.NoError(t, err) + assert.False(t, to11.IsFull()) + assert.False(t, to12.IsFull()) + + assert.True(t, to11.IsEmpty()) + assert.True(t, to12.IsEmpty()) + + stopped := f.Start() + + // write some data + _, errs := fromStep.Write(ctx, writeMessages) + assert.Equal(t, make([]error, 4*batchSize), errs) + + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + otKeys1, _ := otStores["to2"].GetAllKeys(ctx) + loop: + for { + select { + case <-ctx.Done(): + assert.Fail(t, "context cancelled while waiting to get a valid watermark") + break loop + default: + if len(otKeys1) == 0 { + otKeys1, _ = otStores["to1"].GetAllKeys(ctx) + time.Sleep(time.Millisecond * 10) + } else { + // NOTE: in this test we only have one processor to publish + // so len(otKeys) should always be 1 + otKeys1, _ = otStores["to1"].GetAllKeys(ctx) + otValue1, _ := otStores["to1"].GetValue(ctx, otKeys1[0]) + otDecode1, _ := wmb.DecodeToWMB(otValue1) + // if the batch size is 1, then the watermark should be idle because + // one of partitions will be idle + if otDecode1.Idle { + break loop + } else { + time.Sleep(time.Millisecond * 10) + } + } + } + + } + }() + + wg.Add(1) + go func() { + defer wg.Done() + otKeys2, _ := otStores["to2"].GetAllKeys(ctx) + loop: + for { + select { + case <-ctx.Done(): + assert.Fail(t, "context cancelled while waiting to get a valid watermark") + break loop + default: + if len(otKeys2) == 0 { + otKeys2, _ = otStores["to2"].GetAllKeys(ctx) + time.Sleep(time.Millisecond * 10) + } else { + // NOTE: in this test we only have one processor to publish + // so len(otKeys) should always be 1 + otKeys2, _ = otStores["to2"].GetAllKeys(ctx) + otValue2, _ := otStores["to2"].GetValue(ctx, otKeys2[0]) + otDecode2, _ := wmb.DecodeToWMB(otValue2) + // if the batch size is 1, then the watermark should be idle because + // one of partitions will be idle + if otDecode2.Idle { + break loop + } else { + time.Sleep(time.Millisecond * 10) + } + } + } + + } + }() + wg.Wait() + + msgs := to11.GetMessages(1) + for len(msgs) == 0 || msgs[0].Kind != isb.WMB { + select { + case <-ctx.Done(): + logging.FromContext(ctx).Fatalf("expect to have the ctrl message in to1, %s", ctx.Err()) + default: + msgs = to11.GetMessages(1) + time.Sleep(1 * time.Millisecond) + } + } + + msgs = to12.GetMessages(1) + for len(msgs) == 0 || msgs[0].Kind != isb.WMB { + select { + case <-ctx.Done(): + logging.FromContext(ctx).Fatalf("expect to have the ctrl message in to1, %s", ctx.Err()) + default: + msgs = to12.GetMessages(1) + time.Sleep(1 * time.Millisecond) + } + } + + msgs = to21.GetMessages(1) + for len(msgs) == 0 || msgs[0].Kind != isb.WMB { + select { + case <-ctx.Done(): + logging.FromContext(ctx).Fatalf("expect to have the ctrl message in to2, %s", ctx.Err()) + default: + msgs = to21.GetMessages(1) + time.Sleep(1 * time.Millisecond) + } + } + + msgs = to22.GetMessages(1) + for len(msgs) == 0 || msgs[0].Kind != isb.WMB { + select { + case <-ctx.Done(): + logging.FromContext(ctx).Fatalf("expect to have the ctrl message in to2, %s", ctx.Err()) + default: + msgs = to22.GetMessages(1) + time.Sleep(1 * time.Millisecond) + } + } + + // since this is a dropping WhereTo, the buffer can never be full + f.Stop() + + <-stopped + }) + // Explicitly tests the case where we forward to only one buffer + t.Run(tt.name+"_toOneStep", func(t *testing.T) { + batchSize := tt.batchSize + fromStep := simplebuffer.NewInMemoryBuffer("from", 5*batchSize, 0) + to11 := simplebuffer.NewInMemoryBuffer("to1-1", 2*batchSize, 0) + to12 := simplebuffer.NewInMemoryBuffer("to1-2", 2*batchSize, 1) + to21 := simplebuffer.NewInMemoryBuffer("to2-1", 2*batchSize, 0) + to22 := simplebuffer.NewInMemoryBuffer("to2-2", 2*batchSize, 1) + toSteps := map[string][]isb.BufferWriter{ + "to1": {to11, to12}, + "to2": {to21, to22}, + } + ctx, cancel := context.WithTimeout(context.Background(), time.Second*10) + defer cancel() + + writeMessages := testutils.BuildTestWriteMessages(4*batchSize, testStartTime, nil) + + vertex := &dfv1.Vertex{Spec: dfv1.VertexSpec{ + PipelineName: "testPipeline", + AbstractVertex: dfv1.AbstractVertex{ + Name: "testVertex", + }, + }} + + vertexInstance := &dfv1.VertexInstance{ + Vertex: vertex, + Replica: 0, + } + + fetchWatermark := &testForwardFetcher{} + toVertexWmStores := buildWatermarkStores(toSteps) + publishWatermark, otStores := buildPublisherMapAndOTStoreFromWmStores(toSteps, toVertexWmStores) + + defer func() { + for _, p := range publishWatermark { + _ = p.Close() + } + }() + + defer func() { + for _, store := range toVertexWmStores { + _ = store.Close() + } + }() + + idleManager, _ := wmb.NewIdleManager(1, len(toSteps)) + f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, &mySourceForwardTestRoundRobin{}, myForwardTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled)) + + assert.NoError(t, err) + assert.False(t, to11.IsFull()) + assert.False(t, to12.IsFull()) + + assert.True(t, to11.IsEmpty()) + assert.True(t, to12.IsEmpty()) + + stopped := f.Start() + // write some data + _, errs := fromStep.Write(ctx, writeMessages[0:batchSize]) + assert.Equal(t, make([]error, batchSize), errs) + var updatedBatchSize int + if tt.batchSize > 1 { + updatedBatchSize = int(batchSize / 2) + } else { + updatedBatchSize = int(batchSize) + } + // read some data + readMessages, err := to11.Read(ctx, int64(updatedBatchSize)) + assert.NoError(t, err, "expected no error") + + assert.Len(t, readMessages, updatedBatchSize) + for i, j := 0, 0; i < updatedBatchSize; i, j = i+1, j+2 { + assert.Equal(t, []interface{}{writeMessages[j].MessageInfo}, []interface{}{readMessages[i].MessageInfo}) + assert.Equal(t, []interface{}{writeMessages[j].Kind}, []interface{}{readMessages[i].Kind}) + assert.Equal(t, []interface{}{writeMessages[j].Keys}, []interface{}{readMessages[i].Keys}) + assert.Equal(t, []interface{}{writeMessages[j].Body}, []interface{}{readMessages[i].Body}) + } + // write some data + _, errs = fromStep.Write(ctx, writeMessages[batchSize:4*batchSize]) + assert.Equal(t, make([]error, 3*batchSize), errs) + + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + otKeys1, _ := otStores["to1"].GetAllKeys(ctx) + loop: + for { + select { + case <-ctx.Done(): + assert.Fail(t, "context cancelled while waiting to get a valid watermark") + break loop + default: + if len(otKeys1) == 0 { + otKeys1, _ = otStores["to1"].GetAllKeys(ctx) + time.Sleep(time.Millisecond * 10) + } else { + // NOTE: in this test we only have one processor to publish + // so len(otKeys) should always be 1 + otKeys1, _ = otStores["to1"].GetAllKeys(ctx) + otValue1, _ := otStores["to1"].GetValue(ctx, otKeys1[0]) + otDecode1, _ := wmb.DecodeToWMB(otValue1) + if tt.batchSize > 1 && !otDecode1.Idle { + break loop + } else if tt.batchSize == 1 && otDecode1.Idle { + break loop + } else { + time.Sleep(time.Millisecond * 10) + } + } + } + + } + }() + + wg.Add(1) + go func() { + defer wg.Done() + otKeys2, _ := otStores["to2"].GetAllKeys(ctx) + loop: + for { + select { + case <-ctx.Done(): + assert.Fail(t, "context cancelled while waiting to get a valid watermark") + break loop + default: + if len(otKeys2) == 0 { + otKeys2, _ = otStores["to2"].GetAllKeys(ctx) + time.Sleep(time.Millisecond * 10) + } else { + // NOTE: in this test we only have one processor to publish + // so len(otKeys) should always be 1 + otKeys2, _ = otStores["to2"].GetAllKeys(ctx) + otValue2, _ := otStores["to2"].GetValue(ctx, otKeys2[0]) + otDecode2, _ := wmb.DecodeToWMB(otValue2) + // if the batch size is 1, then the watermark should be idle because + // one of partitions will be idle + if otDecode2.Idle { + break loop + } else { + time.Sleep(time.Millisecond * 10) + } + } + } + + } + }() + wg.Wait() + + // stop will cancel the contexts and therefore the forwarder stops without waiting + f.Stop() + + <-stopped + }) + // Test the scenario with UDF error + t.Run(tt.name+"_UDFError", func(t *testing.T) { + batchSize := tt.batchSize + fromStep := simplebuffer.NewInMemoryBuffer("from", 5*batchSize, 0) + to1 := simplebuffer.NewInMemoryBuffer("to1", 2*batchSize, 0) + toSteps := map[string][]isb.BufferWriter{ + "to1": {to1}, + } + + vertex := &dfv1.Vertex{Spec: dfv1.VertexSpec{ + PipelineName: "testPipeline", + AbstractVertex: dfv1.AbstractVertex{ + Name: "testVertex", + }, + }} + + vertexInstance := &dfv1.VertexInstance{ + Vertex: vertex, + Replica: 0, + } + + ctx, cancel := context.WithTimeout(context.Background(), time.Second*10) + defer cancel() + + writeMessages := testutils.BuildTestWriteMessages(4*batchSize, testStartTime, nil) + + fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps) + idleManager, _ := wmb.NewIdleManager(1, len(toSteps)) + f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardApplyUDFErrTest{}, myForwardApplyUDFErrTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled)) + + assert.NoError(t, err) + assert.False(t, to1.IsFull()) + assert.True(t, to1.IsEmpty()) + + stopped := f.Start() + // write some data + _, errs := fromStep.Write(ctx, writeMessages[0:batchSize]) + assert.Equal(t, make([]error, batchSize), errs) + assert.True(t, to1.IsEmpty()) + + f.Stop() + time.Sleep(1 * time.Millisecond) + + <-stopped + }) + // Test the scenario with error + t.Run(tt.name+"_whereToError", func(t *testing.T) { + batchSize := tt.batchSize + fromStep := simplebuffer.NewInMemoryBuffer("from", 5*batchSize, 0) + to1 := simplebuffer.NewInMemoryBuffer("to1", 2*batchSize, 0) + toSteps := map[string][]isb.BufferWriter{ + "to1": {to1}, + } + ctx, cancel := context.WithTimeout(context.Background(), time.Second*10) + defer cancel() + + writeMessages := testutils.BuildTestWriteMessages(4*batchSize, testStartTime, nil) + + vertex := &dfv1.Vertex{Spec: dfv1.VertexSpec{ + PipelineName: "testPipeline", + AbstractVertex: dfv1.AbstractVertex{ + Name: "testVertex", + }, + }} + + vertexInstance := &dfv1.VertexInstance{ + Vertex: vertex, + Replica: 0, + } + + fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps) + idleManager, _ := wmb.NewIdleManager(1, len(toSteps)) + f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardApplyWhereToErrTest{}, myForwardApplyWhereToErrTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled)) + + assert.NoError(t, err) + assert.True(t, to1.IsEmpty()) + + stopped := f.Start() + // write some data + _, errs := fromStep.Write(ctx, writeMessages[0:batchSize]) + assert.Equal(t, make([]error, batchSize), errs) + + f.Stop() + time.Sleep(1 * time.Millisecond) + + assert.True(t, to1.IsEmpty()) + <-stopped + }) + t.Run(tt.name+"_withInternalError", func(t *testing.T) { + batchSize := tt.batchSize + fromStep := simplebuffer.NewInMemoryBuffer("from", 5*batchSize, 0) + to1 := simplebuffer.NewInMemoryBuffer("to1", 2*batchSize, 0) + toSteps := map[string][]isb.BufferWriter{ + "to1": {to1}, + } + ctx, cancel := context.WithTimeout(context.Background(), time.Second*10) + defer cancel() + + writeMessages := testutils.BuildTestWriteMessages(4*batchSize, testStartTime, nil) + + vertex := &dfv1.Vertex{Spec: dfv1.VertexSpec{ + PipelineName: "testPipeline", + AbstractVertex: dfv1.AbstractVertex{ + Name: "testVertex", + }, + }} + + vertexInstance := &dfv1.VertexInstance{ + Vertex: vertex, + Replica: 0, + } + + fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps) + idleManager, _ := wmb.NewIdleManager(1, len(toSteps)) + f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardInternalErrTest{}, myForwardInternalErrTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled)) + + assert.NoError(t, err) + assert.False(t, to1.IsFull()) + assert.True(t, to1.IsEmpty()) + + stopped := f.Start() + // write some data + _, errs := fromStep.Write(ctx, writeMessages[0:batchSize]) + assert.Equal(t, make([]error, batchSize), errs) + + f.Stop() + time.Sleep(1 * time.Millisecond) + <-stopped + }) + } +} + +type testWMBFetcher struct { + // for forward_test.go only + WMBTestSameHeadWMB bool // for testing same head wmb, if set true then WMBTestDiffHeadWMB must be false + sameCounter int + sameLock sync.RWMutex + WMBTestDiffHeadWMB bool // for testing different head wmb, if set true then WMBTestSameHeadWMB must be false + diffCounter int + diffLock sync.RWMutex +} + +// RevertBoolValue set WMBTestSameHeadWMB and WMBTestDiffHeadWMB to opposite value +func (t *testWMBFetcher) RevertBoolValue() { + t.sameLock.Lock() + defer t.sameLock.Unlock() + t.diffLock.Lock() + defer t.diffLock.Unlock() + t.WMBTestSameHeadWMB = !t.WMBTestSameHeadWMB + t.WMBTestDiffHeadWMB = !t.WMBTestDiffHeadWMB +} + +func (t *testWMBFetcher) ComputeWatermark(offset isb.Offset, partition int32) wmb.Watermark { + return t.getWatermark() +} + +// getWatermark uses current time as the watermark because we want to make sure +// the test publisher is publishing watermark +func (t *testWMBFetcher) getWatermark() wmb.Watermark { + return wmb.Watermark(testWMBWatermark) +} + +func (t *testWMBFetcher) ComputeHeadWatermark(int32) wmb.Watermark { + // won't be used + return wmb.Watermark{} +} + +func (t *testWMBFetcher) ComputeHeadIdleWMB(int32) wmb.WMB { + t.sameLock.RLock() + defer t.sameLock.RUnlock() + t.diffLock.RLock() + defer t.diffLock.RUnlock() + if t.WMBTestSameHeadWMB { + t.sameCounter++ + if t.sameCounter == 1 || t.sameCounter == 2 { + return wmb.WMB{ + Idle: true, + Offset: 100, + Watermark: 1636440000000, + } + } + if t.sameCounter == 3 || t.sameCounter == 4 { + return wmb.WMB{ + Idle: true, + Offset: 102, + Watermark: 1636480000000, + } + } + } + if t.WMBTestDiffHeadWMB { + t.diffCounter++ + if t.diffCounter == 1 { + return wmb.WMB{ + Idle: true, + Offset: 100, + Watermark: 1636440000000, + } + } + if t.diffCounter == 2 { + return wmb.WMB{ + Idle: false, + Offset: 101, + Watermark: 1636450000000, + } + } + } + // won't be used + return wmb.WMB{} +} + +func TestNewInterStepDataForwardIdleWatermark(t *testing.T) { + fromStep := simplebuffer.NewInMemoryBuffer("from", 25, 0, simplebuffer.WithReadTimeOut(time.Second)) // default read timeout is 1s + to1 := simplebuffer.NewInMemoryBuffer("to1", 10, 0) + toSteps := map[string][]isb.BufferWriter{ + "to1": {to1}, + } + ctx, cancel := context.WithTimeout(context.Background(), time.Second*10) + defer cancel() + + vertex := &dfv1.Vertex{Spec: dfv1.VertexSpec{ + PipelineName: "testPipeline", + AbstractVertex: dfv1.AbstractVertex{ + Name: "testVertex", + }, + }} + + vertexInstance := &dfv1.VertexInstance{ + Vertex: vertex, + Replica: 0, + } + + ctrlMessage := []isb.Message{{Header: isb.Header{Kind: isb.WMB}}} + writeMessages := testutils.BuildTestWriteMessages(int64(20), testStartTime, nil) + + fetchWatermark := &testWMBFetcher{WMBTestSameHeadWMB: true} + toVertexWmStores := buildWatermarkStores(toSteps) + publishWatermark, otStores := buildPublisherMapAndOTStoreFromWmStores(toSteps, toVertexWmStores) + + defer func() { + for _, p := range publishWatermark { + _ = p.Close() + } + }() + + defer func() { + for _, store := range toVertexWmStores { + _ = store.Close() + } + }() + + idleManager, _ := wmb.NewIdleManager(1, len(toSteps)) + f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardTest{}, myForwardTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(2)) + assert.NoError(t, err) + assert.False(t, to1.IsFull()) + assert.True(t, to1.IsEmpty()) + + stopped := f.Start() + assert.True(t, fromStep.IsEmpty()) + // first batch: read message size is 1 with one ctrl message + // the ctrl message should be acked + // no message published to the next vertex + // so the timeline should be empty + _, errs := fromStep.Write(ctx, ctrlMessage) + assert.Equal(t, make([]error, 1), errs) + // waiting for the ctrl message to be acked + for !fromStep.IsEmpty() { + select { + case <-ctx.Done(): + if ctx.Err() == context.DeadlineExceeded { + t.Fatal("expected the buffer to be empty", ctx.Err()) + } + default: + time.Sleep(1 * time.Millisecond) + } + } + // it should not publish any wm because + // 1. readLength != 0 + // 2. we only have one ctrl message + // 3. meaning dataMessage=0 + otNil, _ := otStores["to1"].GetAllKeys(ctx) + assert.Nil(t, otNil) + + // 2nd and 3rd batches: read message size is 0 + // should send idle watermark + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + otKeys1, _ := otStores["to1"].GetAllKeys(ctx) + for otKeys1 == nil { + otKeys1, _ = otStores["to1"].GetAllKeys(ctx) + time.Sleep(time.Millisecond * 10) + } + }() + wg.Wait() + + otKeys1, _ := otStores["to1"].GetAllKeys(ctx) + otValue1, _ := otStores["to1"].GetValue(ctx, otKeys1[0]) + otDecode1, _ := wmb.DecodeToWMB(otValue1) + for otDecode1.Offset != 0 { // the first ctrl message written to isb. can't use idle because default idle=false + select { + case <-ctx.Done(): + if ctx.Err() == context.DeadlineExceeded { + t.Fatal("expected to have idle watermark in to1 timeline", ctx.Err()) + } + default: + time.Sleep(1 * time.Millisecond) + otKeys1, _ = otStores["to1"].GetAllKeys(ctx) + otValue1, _ = otStores["to1"].GetValue(ctx, otKeys1[0]) + otDecode1, _ = wmb.DecodeToWMB(otValue1) + } + } + assert.Equal(t, wmb.WMB{ + Idle: true, + Offset: 0, // the first ctrl message written to isb + Watermark: 1636440000000, + }, otDecode1) + assert.Equal(t, isb.WMB, to1.GetMessages(1)[0].Kind) + + // 4th batch: read message size = 1 + // a new active watermark should be inserted + _, errs = fromStep.Write(ctx, writeMessages[:2]) + assert.Equal(t, make([]error, 2), errs) + + otKeys1, _ = otStores["to1"].GetAllKeys(ctx) + otValue1, _ = otStores["to1"].GetValue(ctx, otKeys1[0]) + otDecode1, _ = wmb.DecodeToWMB(otValue1) + for otDecode1.Idle { + select { + case <-ctx.Done(): + if ctx.Err() == context.DeadlineExceeded { + t.Fatal("expected to have active watermark in to1 timeline", ctx.Err()) + } + default: + time.Sleep(1 * time.Millisecond) + otKeys1, _ = otStores["to1"].GetAllKeys(ctx) + otValue1, _ = otStores["to1"].GetValue(ctx, otKeys1[0]) + otDecode1, _ = wmb.DecodeToWMB(otValue1) + } + } + assert.Equal(t, wmb.WMB{ + Idle: false, + Offset: 2, // the second message written to isb, read batch size is 2 so the offset is 0+2=2 + Watermark: testWMBWatermark.UnixMilli(), + }, otDecode1) + + // 5th & 6th batch: again idling but got diff head WMB + // so the head is still the same active watermark + // and no new ctrl message to the next vertex + f.wmFetcher.(*testWMBFetcher).RevertBoolValue() + time.Sleep(2 * time.Second) // default read timeout is 1s + otKeys1, _ = otStores["to1"].GetAllKeys(ctx) + otValue1, _ = otStores["to1"].GetValue(ctx, otKeys1[0]) + otDecode1, _ = wmb.DecodeToWMB(otValue1) + assert.Equal(t, wmb.WMB{ + Idle: false, + Offset: 2, // the second message written to isb, read batch size is 2 so the offset is 0+2=2 + Watermark: testWMBWatermark.UnixMilli(), + }, otDecode1) + + var wantKind = []isb.MessageKind{ + isb.WMB, + isb.Data, + isb.Data, + isb.Data, // this is data because it's the default value + } + var wantBody = []bool{ + false, + true, + true, + false, + } + for idx, msg := range to1.GetMessages(4) { + assert.Equal(t, wantKind[idx], msg.Kind) + assert.Equal(t, wantBody[idx], len(msg.Body.Payload) > 0) + } + + // stop will cancel the contexts and therefore the forwarder stops without waiting + f.Stop() + + <-stopped +} + +func TestNewInterStepDataForwardIdleWatermark_Reset(t *testing.T) { + fromStep := simplebuffer.NewInMemoryBuffer("from", 25, 0, simplebuffer.WithReadTimeOut(time.Second)) // default read timeout is 1s + to1 := simplebuffer.NewInMemoryBuffer("to1", 10, 0) + toSteps := map[string][]isb.BufferWriter{ + "to1": {to1}, + } + ctx, cancel := context.WithTimeout(context.Background(), time.Second*10) + defer cancel() + + vertex := &dfv1.Vertex{Spec: dfv1.VertexSpec{ + PipelineName: "testPipeline", + AbstractVertex: dfv1.AbstractVertex{ + Name: "testVertex", + }, + }} + + vertexInstance := &dfv1.VertexInstance{ + Vertex: vertex, + Replica: 0, + } + + writeMessages := testutils.BuildTestWriteMessages(int64(20), testStartTime, nil) + + fetchWatermark := &testWMBFetcher{WMBTestSameHeadWMB: true} + toVertexWmStores := buildWatermarkStores(toSteps) + publishWatermark, otStores := buildPublisherMapAndOTStoreFromWmStores(toSteps, toVertexWmStores) + + defer func() { + for _, p := range publishWatermark { + _ = p.Close() + } + }() + + defer func() { + for _, store := range toVertexWmStores { + _ = store.Close() + } + }() + + // close the fetcher and publishers + defer func() { + for _, p := range publishWatermark { + _ = p.Close() + } + }() + + idleManager, _ := wmb.NewIdleManager(1, len(toSteps)) + f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardTest{}, myForwardTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(2)) + assert.NoError(t, err) + assert.False(t, to1.IsFull()) + assert.True(t, to1.IsEmpty()) + + stopped := f.Start() + var wg sync.WaitGroup + + // 1st and 2nd batches: read message size is 0 + // should send idle watermark + wg.Add(1) + go func() { + defer wg.Done() + otKeys1, _ := otStores["to1"].GetAllKeys(ctx) + for otKeys1 == nil { + otKeys1, _ = otStores["to1"].GetAllKeys(ctx) + time.Sleep(time.Millisecond * 10) + } + }() + wg.Wait() + + otKeys1, _ := otStores["to1"].GetAllKeys(ctx) + otValue1, _ := otStores["to1"].GetValue(ctx, otKeys1[0]) + otDecode1, _ := wmb.DecodeToWMB(otValue1) + for otDecode1.Offset != 0 { // the first ctrl message written to isb. can't use idle because default idle=false + select { + case <-ctx.Done(): + if ctx.Err() == context.DeadlineExceeded { + t.Fatal("expected to have idle watermark in to1 timeline", ctx.Err()) + } + default: + time.Sleep(1 * time.Millisecond) + otKeys1, _ = otStores["to1"].GetAllKeys(ctx) + otValue1, _ = otStores["to1"].GetValue(ctx, otKeys1[0]) + otDecode1, _ = wmb.DecodeToWMB(otValue1) + } + } + assert.Equal(t, wmb.WMB{ + Idle: true, + Offset: 0, // the first ctrl message written to isb + Watermark: 1636440000000, + }, otDecode1) + assert.Equal(t, isb.WMB, to1.GetMessages(1)[0].Kind) + + // 3rd batch: read message size = 1 + // a new active watermark should be inserted + _, errs := fromStep.Write(ctx, writeMessages[:2]) + assert.Equal(t, make([]error, 2), errs) + + otKeys1, _ = otStores["to1"].GetAllKeys(ctx) + otValue1, _ = otStores["to1"].GetValue(ctx, otKeys1[0]) + otDecode1, _ = wmb.DecodeToWMB(otValue1) + for otDecode1.Idle { + select { + case <-ctx.Done(): + if ctx.Err() == context.DeadlineExceeded { + t.Fatal("expected to have active watermark in to1 timeline", ctx.Err()) + } + default: + time.Sleep(1 * time.Millisecond) + otKeys1, _ = otStores["to1"].GetAllKeys(ctx) + otValue1, _ = otStores["to1"].GetValue(ctx, otKeys1[0]) + otDecode1, _ = wmb.DecodeToWMB(otValue1) + } + } + assert.Equal(t, wmb.WMB{ + Idle: false, + Offset: 2, // the second message written to isb, read batch size is 2 so the offset is 0+2=2 + Watermark: testWMBWatermark.UnixMilli(), + }, otDecode1) + + // 5th & 6th batch: again idling should send a new ctrl message + otKeys1, _ = otStores["to1"].GetAllKeys(ctx) + otValue1, _ = otStores["to1"].GetValue(ctx, otKeys1[0]) + otDecode1, _ = wmb.DecodeToWMB(otValue1) + for otDecode1.Offset != 3 { // the second ctrl message written to isb. can't use idle because default idle=false + select { + case <-ctx.Done(): + if ctx.Err() == context.DeadlineExceeded { + t.Fatal("expected to have idle watermark in to1 timeline", ctx.Err()) + } + default: + time.Sleep(1 * time.Millisecond) + otKeys1, _ = otStores["to1"].GetAllKeys(ctx) + otValue1, _ = otStores["to1"].GetValue(ctx, otKeys1[0]) + otDecode1, _ = wmb.DecodeToWMB(otValue1) + } + } + + var wantKind = []isb.MessageKind{ + isb.WMB, + isb.Data, + isb.Data, + isb.WMB, + } + var wantBody = []bool{ + false, + true, + true, + false, + } + for idx, msg := range to1.GetMessages(4) { + assert.Equal(t, wantKind[idx], msg.Kind) + assert.Equal(t, wantBody[idx], len(msg.Body.Payload) > 0) + } + + // stop will cancel the contexts and therefore the forwarder stops without waiting + f.Stop() + + <-stopped +} + +// mySourceForwardTest tests source data transformer by updating message event time, and then verifying new event time and IsLate assignments. +type mySourceForwardTest struct { +} + +func (f mySourceForwardTest) WhereTo(_ []string, _ []string, s string) ([]forwarder.VertexBuffer, error) { + return []forwarder.VertexBuffer{{ + ToVertexName: "to1", + ToVertexPartitionIdx: 0, + }}, nil +} + +type mySourceForwardTestRoundRobin struct { + count int +} + +func (f *mySourceForwardTestRoundRobin) WhereTo(_ []string, _ []string, s string) ([]forwarder.VertexBuffer, error) { + var output = []forwarder.VertexBuffer{{ + ToVertexName: "to1", + ToVertexPartitionIdx: int32(f.count % 2), + }} + f.count++ + return output, nil +} + +func (f mySourceForwardTest) ApplyMap(ctx context.Context, message *isb.ReadMessage) ([]*isb.WriteMessage, error) { + return func(ctx context.Context, readMessage *isb.ReadMessage) ([]*isb.WriteMessage, error) { + _ = ctx + offset := readMessage.ReadOffset + payload := readMessage.Body.Payload + parentPaneInfo := readMessage.MessageInfo + + // apply source data transformer + _ = payload + // copy the payload + result := payload + var key []string + + writeMessage := isb.Message{ + Header: isb.Header{ + MessageInfo: parentPaneInfo, + ID: offset.String(), + Keys: key, + }, + Body: isb.Body{ + Payload: result, + }, + } + return []*isb.WriteMessage{{Message: writeMessage}}, nil + }(ctx, message) +} + +func (f mySourceForwardTest) ApplyMapStream(ctx context.Context, message *isb.ReadMessage, writeMessageCh chan<- isb.WriteMessage) error { + return func(ctx context.Context, readMessage *isb.ReadMessage, writeMessages chan<- isb.WriteMessage) error { + defer close(writeMessages) + + _ = ctx + offset := readMessage.ReadOffset + payload := readMessage.Body.Payload + parentPaneInfo := readMessage.MessageInfo + + // apply source data transformer + _ = payload + // copy the payload + result := payload + var key []string + + writeMessage := isb.Message{ + Header: isb.Header{ + MessageInfo: parentPaneInfo, + ID: offset.String(), + Keys: key, + }, + Body: isb.Body{ + Payload: result, + }, + } + + writeMessages <- isb.WriteMessage{Message: writeMessage} + return nil + }(ctx, message, writeMessageCh) +} + +// TestSourceWatermarkPublisher is a dummy implementation of isb.SourcePublisher interface +type TestSourceWatermarkPublisher struct { +} + +func (p TestSourceWatermarkPublisher) PublishSourceWatermarks([]*isb.ReadMessage) { + // PublishSourceWatermarks is not tested in forwarder_test.go +} + +func TestInterStepDataForwardSinglePartition(t *testing.T) { + fromStep := simplebuffer.NewInMemoryBuffer("from", 25, 0) + to1 := simplebuffer.NewInMemoryBuffer("to1", 10, 0, simplebuffer.WithReadTimeOut(time.Second*10)) + toSteps := map[string][]isb.BufferWriter{ + "to1": {to1}, + } + vertex := &dfv1.Vertex{Spec: dfv1.VertexSpec{ + PipelineName: "testPipeline", + AbstractVertex: dfv1.AbstractVertex{ + Name: "receivingVertex", + }, + }} + + vertexInstance := &dfv1.VertexInstance{ + Vertex: vertex, + Replica: 0, + } + ctx, cancel := context.WithTimeout(context.Background(), time.Second*10) + defer cancel() + + writeMessages := testutils.BuildTestWriteMessages(int64(20), testStartTime, nil) + fetchWatermark := &testForwardFetcher{} + _, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps) + + // create a forwarder + idleManager, _ := wmb.NewIdleManager(1, len(toSteps)) + f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, mySourceForwardTest{}, mySourceForwardTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(5)) + assert.NoError(t, err) + assert.False(t, to1.IsFull()) + assert.True(t, to1.IsEmpty()) + + stopped := f.Start() + count := int64(2) + // write some data + _, errs := fromStep.Write(ctx, writeMessages[0:count]) + assert.Equal(t, make([]error, count), errs) + + // read some data + readMessages, err := to1.Read(ctx, count) + assert.NoError(t, err, "expected no error") + assert.Len(t, readMessages, int(count)) + assert.Equal(t, []interface{}{writeMessages[0].Header.Keys, writeMessages[1].Header.Keys}, []interface{}{readMessages[0].Header.Keys, readMessages[1].Header.Keys}) + assert.Equal(t, []interface{}{"0-0-receivingVertex-0", "1-0-receivingVertex-0"}, []interface{}{readMessages[0].Header.ID, readMessages[1].Header.ID}) + assert.Equal(t, []interface{}{writeMessages[0].Header.Headers, writeMessages[1].Header.Headers}, []interface{}{readMessages[0].Header.Headers, readMessages[1].Header.Headers}) + + f.Stop() + time.Sleep(1 * time.Millisecond) + // only for shutdown will work as from buffer is not empty + f.ForceStop() + <-stopped +} + +func TestInterStepDataForwardMultiplePartition(t *testing.T) { + fromStep := simplebuffer.NewInMemoryBuffer("from", 25, 0) + to11 := simplebuffer.NewInMemoryBuffer("to1-0", 10, 0, simplebuffer.WithReadTimeOut(time.Second*10)) + to12 := simplebuffer.NewInMemoryBuffer("to1-1", 10, 1, simplebuffer.WithReadTimeOut(time.Second*10)) + toSteps := map[string][]isb.BufferWriter{ + "to1": {to11, to12}, + } + vertex := &dfv1.Vertex{Spec: dfv1.VertexSpec{ + PipelineName: "testPipeline", + AbstractVertex: dfv1.AbstractVertex{ + Name: "receivingVertex", + }, + }} + + vertexInstance := &dfv1.VertexInstance{ + Vertex: vertex, + Replica: 0, + } + ctx, cancel := context.WithTimeout(context.Background(), time.Second*10) + defer cancel() + + writeMessages := testutils.BuildTestWriteMessages(int64(20), testStartTime, nil) + fetchWatermark := &testForwardFetcher{} + _, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps) + + // create a forwarder + + idleManager, _ := wmb.NewIdleManager(1, len(toSteps)) + f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, &mySourceForwardTestRoundRobin{}, mySourceForwardTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(5)) + assert.NoError(t, err) + assert.False(t, to11.IsFull()) + assert.False(t, to12.IsFull()) + assert.True(t, to11.IsEmpty()) + assert.True(t, to12.IsEmpty()) + + stopped := f.Start() + count := int64(4) + // write some data + _, errs := fromStep.Write(ctx, writeMessages[0:count]) + assert.Equal(t, make([]error, count), errs) + + time.Sleep(time.Second) + // read some data + // since we have produced four messages, both the partitions should have two messages each.(we write in round-robin fashion) + readMessages, err := to11.Read(ctx, 2) + assert.NoError(t, err, "expected no error") + assert.Len(t, readMessages, 2) + assert.Equal(t, []interface{}{writeMessages[0].Header.Keys, writeMessages[2].Header.Keys}, []interface{}{readMessages[0].Header.Keys, readMessages[1].Header.Keys}) + assert.Equal(t, []interface{}{"0-0-receivingVertex-0", "2-0-receivingVertex-0"}, []interface{}{readMessages[0].Header.ID, readMessages[1].Header.ID}) + assert.Equal(t, []interface{}{writeMessages[0].Header.Headers, writeMessages[2].Header.Headers}, []interface{}{readMessages[0].Header.Headers, readMessages[1].Header.Headers}) + + time.Sleep(time.Second) + + readMessages, err = to12.Read(ctx, 2) + assert.NoError(t, err, "expected no error") + assert.Len(t, readMessages, 2) + assert.Equal(t, []interface{}{writeMessages[1].Header.Keys, writeMessages[3].Header.Keys}, []interface{}{readMessages[0].Header.Keys, readMessages[1].Header.Keys}) + assert.Equal(t, []interface{}{"1-0-receivingVertex-0", "3-0-receivingVertex-0"}, []interface{}{readMessages[0].Header.ID, readMessages[1].Header.ID}) + assert.Equal(t, []interface{}{writeMessages[1].Header.Headers, writeMessages[3].Header.Headers}, []interface{}{readMessages[0].Header.Headers, readMessages[1].Header.Headers}) + + f.Stop() + time.Sleep(1 * time.Millisecond) + // only for shutdown will work as from buffer is not empty + f.ForceStop() + <-stopped +} + +// TestWriteToBuffer tests two BufferFullWritingStrategies: 1. discarding the latest message and 2. retrying writing until context is cancelled. +func TestWriteToBuffer(t *testing.T) { + tests := []struct { + name string + batchSize int64 + strategy dfv1.BufferFullWritingStrategy + streamEnabled bool + throwError bool + }{ + { + name: "test-discard-latest", + batchSize: 10, + strategy: dfv1.DiscardLatest, + streamEnabled: false, + // should not throw any error as we drop messages and finish writing before context is cancelled + throwError: false, + }, + { + name: "test-retry-until-success", + batchSize: 10, + strategy: dfv1.RetryUntilSuccess, + streamEnabled: false, + // should throw context closed error as we keep retrying writing until context is cancelled + throwError: true, + }, + { + name: "test-discard-latest", + batchSize: 1, + strategy: dfv1.DiscardLatest, + streamEnabled: true, + // should not throw any error as we drop messages and finish writing before context is cancelled + throwError: false, + }, + { + name: "test-retry-until-success", + batchSize: 1, + strategy: dfv1.RetryUntilSuccess, + streamEnabled: true, + // should throw context closed error as we keep retrying writing until context is cancelled + throwError: true, + }, + } + for _, value := range tests { + t.Run(value.name, func(t *testing.T) { + fromStep := simplebuffer.NewInMemoryBuffer("from", 5*value.batchSize, 0) + buffer := simplebuffer.NewInMemoryBuffer("to1", value.batchSize, 0, simplebuffer.WithBufferFullWritingStrategy(value.strategy)) + toSteps := map[string][]isb.BufferWriter{ + "to1": {buffer}, + } + ctx, cancel := context.WithTimeout(context.Background(), time.Second*10) + defer cancel() + vertex := &dfv1.Vertex{Spec: dfv1.VertexSpec{ + PipelineName: "testPipeline", + AbstractVertex: dfv1.AbstractVertex{ + Name: "testVertex", + }, + }} + + vertexInstance := &dfv1.VertexInstance{ + Vertex: vertex, + Replica: 0, + } + + fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps) + idleManager, _ := wmb.NewIdleManager(1, len(toSteps)) + f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardTest{}, myForwardTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(value.batchSize), WithUDFStreaming(value.streamEnabled)) + assert.NoError(t, err) + assert.False(t, buffer.IsFull()) + assert.True(t, buffer.IsEmpty()) + + stopped := f.Start() + go func() { + for !buffer.IsFull() { + select { + case <-ctx.Done(): + logging.FromContext(ctx).Fatalf("not full, %s", ctx.Err()) + default: + time.Sleep(1 * time.Millisecond) + } + } + // stop will cancel the context + f.Stop() + }() + + // try to write to buffer after it is full. + var messageToStep = make(map[string][][]isb.Message) + messageToStep["to1"] = make([][]isb.Message, 1) + writeMessages := testutils.BuildTestWriteMessages(4*value.batchSize, testStartTime, nil) + messageToStep["to1"][0] = append(messageToStep["to1"][0], writeMessages[0:value.batchSize+1]...) + _, err = f.writeToBuffers(ctx, messageToStep) + + assert.Equal(t, value.throwError, err != nil) + if value.throwError { + // assert the number of failed messages + assert.True(t, strings.Contains(err.Error(), "with failed messages:1")) + } + <-stopped + }) + } +} + +type myForwardDropTest struct { +} + +func (f myForwardDropTest) WhereTo(_ []string, _ []string, s string) ([]forwarder.VertexBuffer, error) { + return []forwarder.VertexBuffer{}, nil +} + +func (f myForwardDropTest) ApplyMap(ctx context.Context, message *isb.ReadMessage) ([]*isb.WriteMessage, error) { + return testutils.CopyUDFTestApply(ctx, message) +} + +func (f myForwardDropTest) ApplyMapStream(ctx context.Context, message *isb.ReadMessage, writeMessageCh chan<- isb.WriteMessage) error { + return testutils.CopyUDFTestApplyStream(ctx, message, writeMessageCh) +} + +type myForwardToAllTest struct { + count int +} + +func (f *myForwardToAllTest) WhereTo(_ []string, _ []string, s string) ([]forwarder.VertexBuffer, error) { + var output = []forwarder.VertexBuffer{{ + ToVertexName: "to1", + ToVertexPartitionIdx: int32(f.count % 2), + }, + { + ToVertexName: "to2", + ToVertexPartitionIdx: int32(f.count % 2), + }} + f.count++ + return output, nil +} + +func (f *myForwardToAllTest) ApplyMap(ctx context.Context, message *isb.ReadMessage) ([]*isb.WriteMessage, error) { + return testutils.CopyUDFTestApply(ctx, message) +} + +func (f myForwardToAllTest) ApplyMapStream(ctx context.Context, message *isb.ReadMessage, writeMessageCh chan<- isb.WriteMessage) error { + return testutils.CopyUDFTestApplyStream(ctx, message, writeMessageCh) +} + +type myForwardInternalErrTest struct { +} + +func (f myForwardInternalErrTest) WhereTo(_ []string, _ []string, s string) ([]forwarder.VertexBuffer, error) { + return []forwarder.VertexBuffer{{ + ToVertexName: "to1", + ToVertexPartitionIdx: 0, + }}, nil +} + +func (f myForwardInternalErrTest) ApplyMap(_ context.Context, _ *isb.ReadMessage) ([]*isb.WriteMessage, error) { + return nil, &udfapplier.ApplyUDFErr{ + UserUDFErr: false, + InternalErr: struct { + Flag bool + MainCarDown bool + }{Flag: true, MainCarDown: false}, + Message: "InternalErr test", + } +} + +func (f myForwardInternalErrTest) ApplyMapStream(_ context.Context, _ *isb.ReadMessage, writeMessagesCh chan<- isb.WriteMessage) error { + close(writeMessagesCh) + return &udfapplier.ApplyUDFErr{ + UserUDFErr: false, + InternalErr: struct { + Flag bool + MainCarDown bool + }{Flag: true, MainCarDown: false}, + Message: "InternalErr test", + } +} + +type myForwardApplyWhereToErrTest struct { +} + +func (f myForwardApplyWhereToErrTest) WhereTo(_ []string, _ []string, s string) ([]forwarder.VertexBuffer, error) { + return []forwarder.VertexBuffer{{ + ToVertexName: "to1", + ToVertexPartitionIdx: 0, + }}, fmt.Errorf("whereToStep failed") +} + +func (f myForwardApplyWhereToErrTest) ApplyMap(ctx context.Context, message *isb.ReadMessage) ([]*isb.WriteMessage, error) { + return testutils.CopyUDFTestApply(ctx, message) +} + +func (f myForwardApplyWhereToErrTest) ApplyMapStream(ctx context.Context, message *isb.ReadMessage, writeMessageCh chan<- isb.WriteMessage) error { + return testutils.CopyUDFTestApplyStream(ctx, message, writeMessageCh) +} + +type myForwardApplyUDFErrTest struct { +} + +func (f myForwardApplyUDFErrTest) WhereTo(_ []string, _ []string, s string) ([]forwarder.VertexBuffer, error) { + return []forwarder.VertexBuffer{{ + ToVertexName: "to1", + ToVertexPartitionIdx: 0, + }}, nil +} + +func (f myForwardApplyUDFErrTest) ApplyMap(_ context.Context, _ *isb.ReadMessage) ([]*isb.WriteMessage, error) { + return nil, fmt.Errorf("UDF error") +} + +func (f myForwardApplyUDFErrTest) ApplyMapStream(_ context.Context, _ *isb.ReadMessage, writeMessagesCh chan<- isb.WriteMessage) error { + close(writeMessagesCh) + return fmt.Errorf("UDF error") +} + +func validateMetrics(t *testing.T, batchSize int64) { + metadata := ` + # HELP forwarder_data_read_total Total number of Data Messages Read + # TYPE forwarder_data_read_total counter + ` + expected := ` + forwarder_data_read_total{partition_name="from",pipeline="testPipeline",replica="0",vertex="testVertex",vertex_type="MapUDF"} ` + fmt.Sprintf("%f", float64(batchSize)) + ` + ` + + err := testutil.CollectAndCompare(metrics.ReadDataMessagesCount, strings.NewReader(metadata+expected), "forwarder_data_read_total") + if err != nil { + t.Errorf("unexpected collecting result: %v", err) + } + + writeMetadata := ` + # HELP forwarder_write_total Total number of Messages Written + # TYPE forwarder_write_total counter + ` + var writeExpected string + if batchSize > 1 { + writeExpected = ` + forwarder_write_total{partition_name="to1-1",pipeline="testPipeline",replica="0",vertex="testVertex",vertex_type="MapUDF"} ` + fmt.Sprintf("%f", float64(batchSize)/2) + ` + forwarder_write_total{partition_name="to1-2",pipeline="testPipeline",replica="0",vertex="testVertex",vertex_type="MapUDF"} ` + fmt.Sprintf("%f", float64(batchSize)/2) + ` + ` + } else { + writeExpected = ` + forwarder_write_total{partition_name="to1-1",pipeline="testPipeline",replica="0",vertex="testVertex",vertex_type="MapUDF"} ` + fmt.Sprintf("%f", float64(batchSize)) + ` + forwarder_write_total{partition_name="to1-2",pipeline="testPipeline",replica="0",vertex="testVertex",vertex_type="MapUDF"} ` + fmt.Sprintf("%f", float64(0)) + ` + ` + } + + err = testutil.CollectAndCompare(metrics.WriteMessagesCount, strings.NewReader(writeMetadata+writeExpected), "forwarder_write_total") + if err != nil { + t.Errorf("unexpected collecting result:\n%s", err) + } + + ackMetadata := ` + # HELP forwarder_ack_total Total number of Messages Acknowledged + # TYPE forwarder_ack_total counter + ` + ackExpected := ` + forwarder_ack_total{partition_name="from",pipeline="testPipeline",replica="0",vertex="testVertex",vertex_type="MapUDF"} ` + fmt.Sprintf("%d", batchSize) + ` + ` + + err = testutil.CollectAndCompare(metrics.AckMessagesCount, strings.NewReader(ackMetadata+ackExpected), "forwarder_ack_total") + if err != nil { + t.Errorf("unexpected collecting result:\n%s", err) + } +} + +func metricsReset() { + metrics.ReadDataMessagesCount.Reset() + metrics.WriteMessagesCount.Reset() + metrics.AckMessagesCount.Reset() +} + +func buildWatermarkStores(toBuffers map[string][]isb.BufferWriter) map[string]wmstore.WatermarkStore { + var ctx = context.Background() + watermarkStores := make(map[string]wmstore.WatermarkStore) + for key := range toBuffers { + store, _ := wmstore.BuildInmemWatermarkStore(ctx, fmt.Sprintf(publisherKeyspace, key)) + watermarkStores[key] = store + } + return watermarkStores +} + +func buildPublisherMapAndOTStoreFromWmStores(toBuffers map[string][]isb.BufferWriter, wmStores map[string]wmstore.WatermarkStore) (map[string]publish.Publisher, map[string]kvs.KVStorer) { + var ctx = context.Background() + processorEntity := entity.NewProcessorEntity("publisherTestPod") + publishers := make(map[string]publish.Publisher) + otStores := make(map[string]kvs.KVStorer) + + for key, store := range wmStores { + p := publish.NewPublish(ctx, processorEntity, store, int32(len(toBuffers[key])), publish.WithAutoRefreshHeartbeatDisabled(), publish.WithPodHeartbeatRate(1)) + publishers[key] = p + otStores[key] = store.OffsetTimelineStore() + } + return publishers, otStores +} diff --git a/pkg/flatmap/forward/options.go b/pkg/flatmap/forward/options.go new file mode 100644 index 0000000000..42d80848b9 --- /dev/null +++ b/pkg/flatmap/forward/options.go @@ -0,0 +1,92 @@ +/* +Copyright 2022 The Numaproj Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package forward + +import ( + "time" + + "go.uber.org/zap" + + dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1" + "github.com/numaproj/numaflow/pkg/shared/logging" +) + +// options for forwarding the message +type options struct { + // readBatchSize is the default batch size + readBatchSize int64 + // udfConcurrency sets the concurrency for concurrent map UDF processing + udfConcurrency int + // retryInterval is the time.Duration to sleep before retrying + retryInterval time.Duration + // logger is used to pass the logger variable + logger *zap.SugaredLogger + // enableMapUdfStream indicates whether the message streaming is enabled or not for map UDF processing + enableMapUdfStream bool +} + +type Option func(*options) error + +func DefaultOptions() *options { + return &options{ + readBatchSize: dfv1.DefaultReadBatchSize, + udfConcurrency: dfv1.DefaultReadBatchSize, + retryInterval: time.Millisecond, + logger: logging.NewLogger(), + enableMapUdfStream: false, + } +} + +// WithRetryInterval sets the retry interval +func WithRetryInterval(f time.Duration) Option { + return func(o *options) error { + o.retryInterval = time.Duration(f) + return nil + } +} + +// WithReadBatchSize sets the read batch size +func WithReadBatchSize(f int64) Option { + return func(o *options) error { + o.readBatchSize = f + return nil + } +} + +// WithUDFConcurrency sets concurrency for map UDF processing +func WithUDFConcurrency(f int) Option { + return func(o *options) error { + o.udfConcurrency = f + return nil + } +} + +// WithLogger is used to return logger information +func WithLogger(l *zap.SugaredLogger) Option { + return func(o *options) error { + o.logger = l + return nil + } +} + +// WithUDFStreaming sets streaming for map UDF processing +func WithUDFStreaming(f bool) Option { + return func(o *options) error { + o.enableMapUdfStream = f + return nil + } +} diff --git a/pkg/flatmap/forward/shutdown.go b/pkg/flatmap/forward/shutdown.go new file mode 100644 index 0000000000..020c18510f --- /dev/null +++ b/pkg/flatmap/forward/shutdown.go @@ -0,0 +1,74 @@ +/* +Copyright 2022 The Numaproj Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package forward + +import ( + "fmt" + "sync" + "time" +) + +// Shutdown tracks and enforces the shutdown activity. +type Shutdown struct { + startShutdown bool + forceShutdown bool + initiateTime time.Time + shutdownRequestCtr int + rwlock *sync.RWMutex +} + +// IsShuttingDown returns whether we can stop processing. +func (isdf *InterStepDataForward) IsShuttingDown() (bool, error) { + isdf.Shutdown.rwlock.RLock() + defer isdf.Shutdown.rwlock.RUnlock() + + if isdf.Shutdown.forceShutdown || isdf.Shutdown.startShutdown { + return true, nil + } + + return false, nil +} + +func (s *Shutdown) String() string { + s.rwlock.RLock() + defer s.rwlock.RUnlock() + return fmt.Sprintf("startShutdown:%t forceShutdown:%t shutdownRequestCtr:%d initiateTime:%s", + s.startShutdown, s.forceShutdown, s.shutdownRequestCtr, s.initiateTime) +} + +// Stop stops the processing. +func (isdf *InterStepDataForward) Stop() { + isdf.Shutdown.rwlock.Lock() + defer isdf.Shutdown.rwlock.Unlock() + if isdf.Shutdown.initiateTime.IsZero() { + isdf.Shutdown.initiateTime = time.Now() + } + isdf.Shutdown.startShutdown = true + isdf.Shutdown.shutdownRequestCtr++ + // call cancel + isdf.cancelFn() +} + +// ForceStop sets up the force shutdown flag. +func (isdf *InterStepDataForward) ForceStop() { + // call stop (what if we have an enthusiastic shutdown that forces first) + // e.g. I know I have written a wrong map UDF, so shutdown ASAP + isdf.Stop() + isdf.Shutdown.rwlock.Lock() + defer isdf.Shutdown.rwlock.Unlock() + isdf.Shutdown.forceShutdown = true +} diff --git a/pkg/flatmap/forward/shutdown_test.go b/pkg/flatmap/forward/shutdown_test.go new file mode 100644 index 0000000000..f6fb478d36 --- /dev/null +++ b/pkg/flatmap/forward/shutdown_test.go @@ -0,0 +1,156 @@ +/* +Copyright 2022 The Numaproj Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package forward + +import ( + "context" + "testing" + "time" + + dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1" + "github.com/numaproj/numaflow/pkg/forwarder" + "github.com/numaproj/numaflow/pkg/isb" + "github.com/numaproj/numaflow/pkg/isb/stores/simplebuffer" + "github.com/numaproj/numaflow/pkg/isb/testutils" + "github.com/numaproj/numaflow/pkg/watermark/generic" + "github.com/numaproj/numaflow/pkg/watermark/wmb" + + "github.com/stretchr/testify/assert" +) + +type myShutdownTest struct { +} + +func (s myShutdownTest) WhereTo(_ []string, _ []string, _ string) ([]forwarder.VertexBuffer, error) { + return []forwarder.VertexBuffer{}, nil +} + +func (s myShutdownTest) ApplyMap(ctx context.Context, message *isb.ReadMessage) ([]*isb.WriteMessage, error) { + return testutils.CopyUDFTestApply(ctx, message) +} + +func (s myShutdownTest) ApplyMapStream(ctx context.Context, message *isb.ReadMessage, writeMessageCh chan<- isb.WriteMessage) error { + return testutils.CopyUDFTestApplyStream(ctx, message, writeMessageCh) +} + +func TestInterStepDataForward(t *testing.T) { + tests := []struct { + name string + batchSize int64 + streamEnabled bool + }{ + { + name: "stream_forward", + batchSize: 1, + streamEnabled: true, + }, + { + name: "batch_forward", + batchSize: 5, + streamEnabled: false, + }, + } + for _, tt := range tests { + t.Run(tt.name+"_stop", func(t *testing.T) { + batchSize := tt.batchSize + fromStep := simplebuffer.NewInMemoryBuffer("from", 5*batchSize, 0) + to1 := simplebuffer.NewInMemoryBuffer("to1", 2*batchSize, 0) + toSteps := map[string][]isb.BufferWriter{ + "to1": {to1}, + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + startTime := time.Unix(1636470000, 0) + writeMessages := testutils.BuildTestWriteMessages(4*batchSize, startTime, nil) + + vertex := &dfv1.Vertex{Spec: dfv1.VertexSpec{ + PipelineName: "testPipeline", + AbstractVertex: dfv1.AbstractVertex{ + Name: "testVertex", + }, + }} + + vertexInstance := &dfv1.VertexInstance{ + Vertex: vertex, + Replica: 0, + } + + fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps) + + idleManager, _ := wmb.NewIdleManager(1, len(toSteps)) + f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myShutdownTest{}, myShutdownTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled)) + assert.NoError(t, err) + stopped := f.Start() + // write some data but buffer is not full even though we are not reading + _, errs := fromStep.Write(ctx, writeMessages[0:batchSize]) + assert.Equal(t, make([]error, batchSize), errs) + + f.Stop() + // we cannot assert the result of IsShuttingDown because it might take a couple of iterations to be successful. + _, _ = f.IsShuttingDown() + <-stopped + }) + t.Run(tt.name+"_forceStop", func(t *testing.T) { + batchSize := tt.batchSize + fromStep := simplebuffer.NewInMemoryBuffer("from", 5*batchSize, 0) + to1 := simplebuffer.NewInMemoryBuffer("to", 2*batchSize, 0) + toSteps := map[string][]isb.BufferWriter{ + "to1": {to1}, + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + startTime := time.Unix(1636470000, 0) + writeMessages := testutils.BuildTestWriteMessages(4*batchSize, startTime, nil) + + vertex := &dfv1.Vertex{Spec: dfv1.VertexSpec{ + PipelineName: "testPipeline", + AbstractVertex: dfv1.AbstractVertex{ + Name: "testVertex", + }, + }} + + vertexInstance := &dfv1.VertexInstance{ + Vertex: vertex, + Replica: 0, + } + + fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps) + + idleManager, _ := wmb.NewIdleManager(1, len(toSteps)) + f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myShutdownTest{}, myShutdownTest{}, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled)) + assert.NoError(t, err) + stopped := f.Start() + // write some data such that the fromBufferPartition can be empty, that is toBuffer gets full + _, errs := fromStep.Write(ctx, writeMessages[0:4*batchSize]) + assert.Equal(t, make([]error, 4*batchSize), errs) + + f.Stop() + canIShutdown, _ := f.IsShuttingDown() + assert.Equal(t, true, canIShutdown) + time.Sleep(1 * time.Millisecond) + // only for canIShutdown will work as from buffer is not empty + f.ForceStop() + canIShutdown, err = f.IsShuttingDown() + assert.NoError(t, err) + assert.Equal(t, true, canIShutdown) + + <-stopped + }) + } +} diff --git a/pkg/flatmap/rpc/errors.go b/pkg/flatmap/rpc/errors.go new file mode 100644 index 0000000000..cfdffe5757 --- /dev/null +++ b/pkg/flatmap/rpc/errors.go @@ -0,0 +1,51 @@ +/* +Copyright 2022 The Numaproj Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package rpc + +import "fmt" + +// ApplyUDFErr represents any mapUDF related error +type ApplyUDFErr struct { + UserUDFErr bool + Message string + InternalErr +} + +// InternalErr represents errors internal to the platform +type InternalErr struct { + Flag bool + MainCarDown bool +} + +// IsUserUDFErr is true if the problem is due to the user code in the UDF. +func (e *ApplyUDFErr) IsUserUDFErr() bool { + return e.UserUDFErr +} + +// IsInternalErr is true if this is a platform issue. This is a blocking error. +func (e *ApplyUDFErr) IsInternalErr() bool { + return e.InternalErr.Flag +} + +func (e *ApplyUDFErr) Error() string { + return fmt.Sprint(e.Message) +} + +// Is checks if the error is of the same type +func (e *ApplyUDFErr) Is(target error) bool { + return target.Error() == e.Error() +} diff --git a/pkg/flatmap/rpc/grpc_flatmap.go b/pkg/flatmap/rpc/grpc_flatmap.go new file mode 100644 index 0000000000..28bf8824b4 --- /dev/null +++ b/pkg/flatmap/rpc/grpc_flatmap.go @@ -0,0 +1,270 @@ +package rpc + +import ( + context2 "context" + "fmt" + "time" + + flatmappb "github.com/numaproj/numaflow-go/pkg/apis/proto/flatmap/v1" + "golang.org/x/net/context" + "google.golang.org/protobuf/types/known/emptypb" + "google.golang.org/protobuf/types/known/timestamppb" + + "github.com/numaproj/numaflow/pkg/flatmap/tracker" + "github.com/numaproj/numaflow/pkg/flatmap/types" + sdkerr "github.com/numaproj/numaflow/pkg/sdkclient/error" + "github.com/numaproj/numaflow/pkg/sdkclient/flatmapper" + "github.com/numaproj/numaflow/pkg/shared/logging" +) + +// GRPCBasedFlatmap is a flat map applier that uses gRPC client to invoke the flat map UDF. +// It implements the applier.FlatmapApplier interface. +type GRPCBasedFlatmap struct { + client flatmapper.Client + tracker *tracker.Tracker + vertexName string +} + +func NewUDSgRPCBasedFlatmap(client flatmapper.Client, vertexName string) *GRPCBasedFlatmap { + return &GRPCBasedFlatmap{client: client, tracker: tracker.NewTracker(), vertexName: vertexName} +} + +// IsHealthy checks if the map udf is healthy. +func (u *GRPCBasedFlatmap) IsHealthy(ctx context.Context) error { + return u.WaitUntilReady(ctx) +} + +// CloseConn closes the gRPC client connection. +func (u *GRPCBasedFlatmap) CloseConn(ctx context.Context) error { + return u.client.CloseConn(ctx) +} + +// WaitUntilReady waits until the reduce udf is connected. +func (u *GRPCBasedFlatmap) WaitUntilReady(ctx context.Context) error { + logger := logging.FromContext(ctx) + for { + select { + case <-ctx.Done(): + return fmt.Errorf("failed on readiness check: %w", ctx.Err()) + default: + if _, err := u.client.IsReady(ctx, &emptypb.Empty{}); err == nil { + return nil + } else { + logger.Infof("waiting for reduce udf to be ready: %v", err) + time.Sleep(1 * time.Second) + } + } + } +} + +// ApplyMap applies the map udf on the stream of read messages and streams the responses back on the responseCh +// Internally, it spawns two go-routines, one for sending the requests to the client and the other to listen to the +// responses back from it. +func (u *GRPCBasedFlatmap) ApplyMap(ctx context2.Context, messageStream []*types.RequestFlatmap, responseCh chan<- *flatmappb.MapResponse) (chan struct{}, <-chan error) { + // errCh is used to propagate any errors recieved from the grpc client upstream so that they can be handled + // accordingly. + errCh := make(chan error) + doneChan := make(chan struct{}) + + // flatmapRequests is a channel on which the input requests are streamed, this is then consumed by the grpc client + //TODO(stream): do we need to keep this buffered? + flatmapRequests := make(chan *flatmappb.MapRequest) + + // Response routine: + // This routine would invoke the MapFn from the client and then keep listening to the response and errorCh + // from the same. + // On getting a response, it would parse it to check whether this is the last response for a given request, + // in such a case we will remove it from the tracker. + // It would also look for any errors received from the client, and then propagate them. + // TODO(stream): should we move this tracking mechanism on a higher layer, and track a request till the + // end of lifetime ie ack + go func() { + // TODO(stream): Instead of closing the channel here, return a done and close this upstream? + // close the responseCh while exiting to indicate downstream that no more responses expected from + // gRPC + defer close(doneChan) + // invoke the MapFn from the gRPC client for a stream of input requests + // resultCh -> chan to read responses streamed back + // mapErrCh -> chan for reading any errors encountered during gRPC + resultCh, mapErrCh := u.client.MapFn(ctx, flatmapRequests) + // Keep running forever until explicit return + for { + // See if we got a response from the client, could be on the response or the error channel + select { + // Got a response on the resultCh + case result, ok := <-resultCh: + // If there are no more messages to read on the stream, or a nil message we can safely assume that + // gRPC has no more messages to send. Hence, we can return from here + if !ok || result == nil { + return + } + //resp, remove, uid := u.parseMapResponse(result) + //// If this was the last response for a request, let's remove from the tracker + //// As this is a special message indicates that all responses for a request have been + //// received, the ackIt field is enabled to indicate that the parent request can be acked now. + //if remove { + // u.tracker.RemoveRequest(uid) + //} + // Forward the received response to the channel + //log.Print("MYDEBUG SENDING TO WRITER ", result.Result.GetUuid()) + responseCh <- result + case err := <-mapErrCh: + // We got a context done while processing the gRPC, hence stop processing + // The specific case for ctx.Done() is already checked in MapFn + if err == ctx.Err() { + errCh <- err + return + } + // If we got any other type of error + if err != nil { + // TODO(stream): graceful handling of error, so that we can drain all the + // unprocessed messaged and retry them again. Once way could be to just restart the NUMA container + // in such a case, which would force a reread of the messages which have not been acked. + errCh <- convertToUdfError(err) + // TODO(stream): Should we stop processing further in this case then + return + } + } + } + }() + + // Read routine: this goroutine reads on the messageStream slice and sends each + // of the read messages to the grpc client + // after transforming it to a MapRequest. Once all messages are sent, it closes the input channel + // to indicate that all requests have been read. + // On creating a new request, we add it to a tracker map so that the responses on the stream + // can be mapped backed to the given parent request + go func() { + defer close(flatmapRequests) + for _, req := range messageStream { + d := u.createMapRequest(req) + //log.Print("MYDEBUG SENDING TO client ", d.GetUuid()) + flatmapRequests <- d + } + }() + return doneChan, errCh +} + +// createMapRequest takes an isb.ReadMessage and returns proto MapRequest +func (u *GRPCBasedFlatmap) createMapRequest(reqMsg *types.RequestFlatmap) *flatmappb.MapRequest { + msg := reqMsg.Request + keys := msg.Keys + payload := msg.Body.Payload + //// Add the request to the tracker, and get the unique UUID corresponding to it + //uid := u.tracker.AddRequest(msg) + // Create the MapRequest, with the required fields. + var d = &flatmappb.MapRequest{ + Keys: keys, + Value: payload, + EventTime: timestamppb.New(msg.EventTime), + Watermark: timestamppb.New(msg.Watermark), + Headers: msg.Headers, + Uuid: reqMsg.Uid, + } + return d +} + +// parseMapResponse takes a proto response from the gRPC and converts this into a ResponseFlatmap type, +// this also checks if this was a special EOR response, in such a case we indicate that the request corresponding +// to the response can be safely removed from the tracker. +func (u *GRPCBasedFlatmap) parseMapResponse(resp *flatmappb.MapResponse) (parsedResp *types.ResponseFlatmap, requestDone bool, uid string) { + result := resp.Result + eor := result.GetEOR() + uid = result.GetUuid() + //parentRequest, ok := u.tracker.GetRequest(uid) + // TODO(stream): check what should be path for !ok, which means that we got a UUID + // which has already been deleted from the tracker/ or never added in the first place + // can this even happen though if messages are ordered and we only have a single routine processing it? + //if !ok { + //} + // Request has completed remove from the tracker module + if eor == true { + return &types.ResponseFlatmap{ + //ParentMessage: parentRequest, + Uid: uid, + RespMessage: nil, + AckIt: true, + Total: int64(resp.Result.Total), + }, true, uid + } + //keys := result.GetKeys() + //taggedMessage := &isb.WriteMessage{ + // Message: isb.Message{ + // Header: isb.Header{ + // MessageInfo: parentRequest.MessageInfo, + // // We need this to be unique so that the ISB can execute its Dedup logic + // // this ID should be such that even when the same response is processed and received + // // again from the UDF, we still assign it the same ID. + // // The ID here will be a concat of the three values + // // parentRequest.ReadOffset - vertexName - result.Index + // // + // // ReadOffset - Will be the read offset of the request which corresponds to this response. + // // We have this stored in our tracker. + // // + // // VertexName - the name of the vertex from which this response is generated, this is + // // important to ensure that we can differentiate between messages emitted from 2 map vertices + // // + // // Result Index - This parameter is added on the SDK side. + // // We add the index of the message from the messages slice to the individual response. + // // TODO(stream): explore if there can be more robust ways to do this + // ID: fmt.Sprintf("%s-%s-%s", parentRequest.ReadOffset.String(), u.vertexName, result.GetIndex()), + // Keys: keys, + // }, + // Body: isb.Body{ + // Payload: result.GetValue(), + // }, + // }, + // Tags: result.GetTags(), + //} + return &types.ResponseFlatmap{ + //ParentMessage: parentRequest, + Uid: uid, + //RespMessage: result, + AckIt: false, + Total: int64(resp.Result.Total), + }, false, uid +} + +// convertToUdfError converts the error returned by the MapFn to ApplyUDFErr +func convertToUdfError(err error) error { + udfErr, _ := sdkerr.FromError(err) + switch udfErr.ErrorKind() { + case sdkerr.Retryable: + // TODO: currently we don't handle retryable errors yet + return &ApplyUDFErr{ + UserUDFErr: false, + Message: fmt.Sprintf("gRPC client.MapFn failed, %s", err), + InternalErr: InternalErr{ + Flag: true, + MainCarDown: false, + }, + } + case sdkerr.NonRetryable: + return &ApplyUDFErr{ + UserUDFErr: false, + Message: fmt.Sprintf("gRPC client.MapFn failed, %s", err), + InternalErr: InternalErr{ + Flag: true, + MainCarDown: false, + }, + } + case sdkerr.Canceled: + return &ApplyUDFErr{ + UserUDFErr: false, + Message: context.Canceled.Error(), + InternalErr: InternalErr{ + Flag: true, + MainCarDown: false, + }, + } + default: + return &ApplyUDFErr{ + UserUDFErr: false, + Message: fmt.Sprintf("gRPC client.MapFn failed, %s", err), + InternalErr: InternalErr{ + Flag: true, + MainCarDown: false, + }, + } + } +} diff --git a/pkg/flatmap/tracker/tracker.go b/pkg/flatmap/tracker/tracker.go new file mode 100644 index 0000000000..2512284efe --- /dev/null +++ b/pkg/flatmap/tracker/tracker.go @@ -0,0 +1,128 @@ +package tracker + +import ( + "log" + "sync" + + "github.com/google/uuid" + + "github.com/numaproj/numaflow/pkg/flatmap/types" + "github.com/numaproj/numaflow/pkg/isb" +) + +type Tracker struct { + requestMap sync.Map + // TODO(stream): check perf between sync.map and mutex+map + responseIdx sync.Map + lock sync.RWMutex + m map[string]*types.RequestFlatmap +} + +func NewTracker() *Tracker { + return &Tracker{ + //requestMap: sync.Map{}, + //responseIdx: sync.Map{}, + m: make(map[string]*types.RequestFlatmap), + } +} + +func GetNewId() string { + id, _ := uuid.NewUUID() + return id.String() +} + +//func (t *Tracker) AddRequest(msg *isb.ReadMessage) string { +// // TODO(stream): we could use read offset as the ID now instead of UUID? +// id := GetNewId() +// //t.requestMap.Store(id, msg) +// t.Set(id, msg) +// return id +//} + +func (t *Tracker) AddRequest(msg *isb.ReadMessage) *types.RequestFlatmap { + // TODO(stream): we could use read offset as the ID now instead of UUID? + id := GetNewId() + flatmapRequest := &types.RequestFlatmap{ + Request: msg, + Uid: id, + ReadOffset: msg.ReadOffset, + } + //t.requestMap.Store(id, msg) + t.Set(id, flatmapRequest) + return flatmapRequest +} + +func (t *Tracker) GetRequest(id string) (*types.RequestFlatmap, bool) { + return t.Get(id) +} + +func (t *Tracker) NewResponse(id string) { + t.responseIdx.Store(id, 1) +} + +func (t *Tracker) IncrementRespIdx(id string) bool { + idx, ok := t.responseIdx.Load(id) + if !ok { + return ok + } + newIdx := idx.(int) + 1 + t.responseIdx.Store(id, newIdx) + return true +} + +func (t *Tracker) GetIdx(id string) (int, bool) { + idx, ok := t.responseIdx.Load(id) + if !ok { + return -1, ok + } + return idx.(int), ok +} + +func (t *Tracker) RemoveRequest(id string) { + //t.requestMap.Delete(id) + //t.responseIdx.Delete(id) + t.Delete(id) +} + +func (t *Tracker) PrintAll() { + t.lock.RLock() + defer t.lock.RUnlock() + for k, v := range t.m { + log.Println("MYDEBUG: MAP VALS", k, " ", v) + } +} + +func (t *Tracker) Get(key string) (*types.RequestFlatmap, bool) { + t.lock.RLock() + defer t.lock.RUnlock() + item, ok := t.m[key] + return item, ok +} + +func (t *Tracker) Set(key string, value *types.RequestFlatmap) { + t.lock.Lock() + defer t.lock.Unlock() + t.m[key] = value +} + +func (t *Tracker) Delete(key string) { + t.lock.Lock() + defer t.lock.Unlock() + delete(t.m, key) +} + +func (t *Tracker) GetItems() []*isb.ReadMessage { + t.lock.Lock() + defer t.lock.Unlock() + items := make([]*isb.ReadMessage, 0, len(t.m)) + for _, vals := range t.m { + items = append(items, vals.Request) + } + return items +} + +func (t *Tracker) Clear() { + t.lock.Lock() + defer t.lock.Unlock() + t.m = make(map[string]*types.RequestFlatmap) +} diff --git a/pkg/flatmap/types/types.go b/pkg/flatmap/types/types.go new file mode 100644 index 0000000000..7db6bd17b1 --- /dev/null +++ b/pkg/flatmap/types/types.go @@ -0,0 +1,33 @@ +package types + +import ( + flatmappb "github.com/numaproj/numaflow-go/pkg/apis/proto/flatmap/v1" + + "github.com/numaproj/numaflow/pkg/isb" +) + +// TODO(stream): check what all data to keep here, this might be a lot to keep the whole message? +// We might just need only a few details out of these +type ResponseFlatmap struct { + ParentMessage *isb.ReadMessage + Uid string + RespMessage *flatmappb.MapResponse + AckIt bool + Total int64 +} + +type WriteMsgFlatmap struct { + Message *ResponseFlatmap + AckIt bool +} + +type AckMsgFlatmap struct { + Message *isb.ReadMessage + AckIt bool +} + +type RequestFlatmap struct { + Request *isb.ReadMessage + Uid string + ReadOffset isb.Offset +} diff --git a/pkg/isb/interfaces.go b/pkg/isb/interfaces.go index 6593ec0826..7a331f282b 100644 --- a/pkg/isb/interfaces.go +++ b/pkg/isb/interfaces.go @@ -35,6 +35,7 @@ type BufferWriter interface { BufferWriterInformation io.Closer Write(context.Context, []Message) ([]Offset, []error) + WriteNew(context.Context, Message) (Offset, error) } // BufferReader is the buffer from which we are reading. diff --git a/pkg/isb/stores/jetstream/reader.go b/pkg/isb/stores/jetstream/reader.go index 72e643e9cc..f2db09964c 100644 --- a/pkg/isb/stores/jetstream/reader.go +++ b/pkg/isb/stores/jetstream/reader.go @@ -152,6 +152,19 @@ func (jr *jetStreamReader) Ack(_ context.Context, offsets []isb.Offset) []error defer func(t time.Time) { isbAckTime.With(labels).Observe(float64(time.Since(t).Microseconds())) }(time.Now()) + + if len(offsets) == 1 { + if err := offsets[0].AckIt(); err != nil { + jr.log.Errorw("Failed to ack message", zap.Error(err)) + // If the error is related to nats/jetstream, we skip it because it might end up with infinite ack retries. + // Skipping those errors to let the whole read/write/ack for loop to restart from reading, to pick up those + // redelivered messages. + if !strings.HasPrefix(err.Error(), "nats:") { + return []error{err} + } + } + return nil + } errs := make([]error, len(offsets)) done := make(chan struct{}) wg := &sync.WaitGroup{} diff --git a/pkg/isb/stores/jetstream/writer.go b/pkg/isb/stores/jetstream/writer.go index 8605b3290d..e23b7d35cd 100644 --- a/pkg/isb/stores/jetstream/writer.go +++ b/pkg/isb/stores/jetstream/writer.go @@ -151,6 +151,34 @@ func (jw *jetStreamWriter) Close() error { return nil } +func (jw *jetStreamWriter) WriteNew(ctx context.Context, message isb.Message) (isb.Offset, error) { + labels := map[string]string{"buffer": jw.GetName()} + errs := fmt.Errorf("unknown error") + if jw.isFull.Load() { + jw.log.Debugw("Is full") + isbFull.With(labels).Inc() + // when buffer is full, we need to decide whether to discard the message or not. + switch jw.opts.bufferFullWritingStrategy { + case v1alpha1.DiscardLatest: + // user explicitly wants to discard the message when buffer if full. + // return no retryable error as a callback to let caller know that the message is discarded. + errs = isb.NonRetryableBufferWriteErr{Name: jw.name, Message: isb.BufferFullMessage} + default: + // Default behavior is to return a BufferWriteErr. + errs = isb.BufferWriteErr{Name: jw.name, Full: true, Message: isb.BufferFullMessage} + } + isbWriteErrors.With(labels).Inc() + return nil, errs + } + //// TODO: temp env flag for sync/async writing, revisit this later. + //if sharedutil.LookupEnvStringOr("ISB_ASYNC_WRITE", "false") == "true" { + // return jw.asyncWrite(ctx, messages, errs, labels) + //} + + routine, err := jw.WriteRoutine(message, labels) + return routine, err +} + func (jw *jetStreamWriter) Write(ctx context.Context, messages []isb.Message) ([]isb.Offset, []error) { labels := map[string]string{"buffer": jw.GetName()} var errs = make([]error, len(messages)) @@ -256,10 +284,48 @@ func (jw *jetStreamWriter) asyncWrite(_ context.Context, messages []isb.Message, return writeOffsets, errs } +func (jw *jetStreamWriter) WriteRoutine(message isb.Message, metricsLabels map[string]string) (isb.Offset, error) { + defer func(t time.Time) { + isbWriteTime.With(metricsLabels).Observe(float64(time.Since(t).Microseconds())) + }(time.Now()) + payload, err := message.MarshalBinary() + if err != nil { + return nil, err + } + m := &nats.Msg{ + Subject: jw.subject, + Data: payload, + } + pubOpts := []nats.PubOpt{nats.AckWait(2 * time.Second)} + // nats.MsgId() is for exactly-once writing + // we don't need to set MsgId for control message + if message.Header.Kind != isb.WMB { + pubOpts = append(pubOpts, nats.MsgId(message.Header.ID)) + } + if pubAck, err := jw.js.PublishMsg(m, pubOpts...); err != nil { + isbWriteErrors.With(metricsLabels).Inc() + return nil, err + } else { + if pubAck.Duplicate { + // If a message gets repeated, it will have the same offset number as the one before it. + // We shouldn't try to publish watermark on these repeated messages. Doing so would + // violate the principle of publishing watermarks to monotonically increasing offsets. + return nil, isb.NonRetryableBufferWriteErr{Name: jw.name, Message: isb.DuplicateIDMessage} + } else { + jw.log.Debugw("Succeeded to publish a message", zap.String("stream", pubAck.Stream), zap.Any("seq", pubAck.Sequence), zap.Bool("duplicate", pubAck.Duplicate), zap.String("msgID", message.Header.ID), zap.String("domain", pubAck.Domain)) + return &writeOffset{seq: pubAck.Sequence, partitionIdx: jw.partitionIdx}, nil + } + } +} + func (jw *jetStreamWriter) syncWrite(_ context.Context, messages []isb.Message, errs []error, metricsLabels map[string]string) ([]isb.Offset, []error) { defer func(t time.Time) { isbWriteTime.With(metricsLabels).Observe(float64(time.Since(t).Microseconds())) }(time.Now()) + if len(messages) == 1 { + routine, err := jw.WriteRoutine(messages[0], metricsLabels) + return []isb.Offset{routine}, []error{err} + } var writeOffsets = make([]isb.Offset, len(messages)) wg := new(sync.WaitGroup) for index, msg := range messages { diff --git a/pkg/isb/stores/redis/write.go b/pkg/isb/stores/redis/write.go index 94876feffc..640e11d759 100644 --- a/pkg/isb/stores/redis/write.go +++ b/pkg/isb/stores/redis/write.go @@ -51,6 +51,11 @@ type BufferWrite struct { log *zap.SugaredLogger } +func (bw *BufferWrite) WriteNew(ctx context.Context, message isb.Message) (isb.Offset, error) { + //TODO implement me + panic("implement me") +} + // BufferWriteInfo will contain the buffer infoRefreshInterval from the writer point of view. type BufferWriteInfo struct { isFull *atomic.Bool diff --git a/pkg/isb/stores/simplebuffer/buffer.go b/pkg/isb/stores/simplebuffer/buffer.go index 893c2fe085..99832dbe36 100644 --- a/pkg/isb/stores/simplebuffer/buffer.go +++ b/pkg/isb/stores/simplebuffer/buffer.go @@ -45,6 +45,11 @@ type InMemoryBuffer struct { rwlock *sync.RWMutex } +func (b *InMemoryBuffer) WriteNew(ctx context.Context, message isb.Message) (isb.Offset, error) { + //TODO implement me + panic("implement me") +} + var _ isb.BufferReader = (*InMemoryBuffer)(nil) var _ isb.BufferWriter = (*InMemoryBuffer)(nil) diff --git a/pkg/sdkclient/const.go b/pkg/sdkclient/const.go index 4b644d4b7e..d1f76fd225 100644 --- a/pkg/sdkclient/const.go +++ b/pkg/sdkclient/const.go @@ -34,8 +34,10 @@ const ( SourceAddr = "/var/run/numaflow/source.sock" SourceTransformerAddr = "/var/run/numaflow/sourcetransform.sock" MultiProcAddr = "/var/run/numaflow/multiproc" + FlatmapAddr = "/var/run/numaflow/flatmap.sock" // Server information file configs + FlatmapServerInfoFile = "/var/run/numaflow/flatmapper-server-info" MapServerInfoFile = "/var/run/numaflow/mapper-server-info" MapStreamServerInfoFile = "/var/run/numaflow/mapstreamer-server-info" ReduceServerInfoFile = "/var/run/numaflow/reducer-server-info" diff --git a/pkg/sdkclient/flatmapper/client.go b/pkg/sdkclient/flatmapper/client.go new file mode 100644 index 0000000000..cf549d3a6d --- /dev/null +++ b/pkg/sdkclient/flatmapper/client.go @@ -0,0 +1,235 @@ +/* +Copyright 2022 The Numaproj Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package flatmapper + +import ( + "context" + "errors" + "io" + "strconv" + "strings" + + flatmappb "github.com/numaproj/numaflow-go/pkg/apis/proto/flatmap/v1" + "github.com/numaproj/numaflow-go/pkg/info" + "google.golang.org/grpc" + "google.golang.org/protobuf/types/known/emptypb" + + "github.com/numaproj/numaflow/pkg/sdkclient" + sdkerr "github.com/numaproj/numaflow/pkg/sdkclient/error" + grpcutil "github.com/numaproj/numaflow/pkg/sdkclient/grpc" +) + +// client contains the grpc connection and the grpc client. +type client struct { + conn *grpc.ClientConn + grpcClt flatmappb.FlatmapClient +} + +// MapFn is the RPC handler for the gRPC client (Numa container) +// It takes in a stream of input Requests, sends them to the gRPC server(UDF) and then streams the +// responses received back on a channel asynchronously. +// We spawn 2 goroutines here, one for sending the requests over the stream +// and the other one for receiving the responses +func (c client) MapFn(ctx context.Context, datumStreamCh <-chan *flatmappb.MapRequest) (<-chan *flatmappb.MapResponse, <-chan error) { + var ( + // errCh is used to track and propagate any errors that might occur during the rpc lifecyle, these could include + // errors in sending, UDF errors etc + // These are propagated to the applier for further handling + errCh = make(chan error) + // TODO(stream): Should we keep this buffered? Might help with error scenario to drain any + // messages already processed + responseCh = make(chan *flatmappb.MapResponse) + ) + + //MapFn is a bidirectional RPC + //We get a Flatmap_MapFnClient interface over which we can send the requests, + //receive the responses asynchronously. + //TODO(stream): this creates a new gRPC stream for every batch, + //it might be useful to see the performance difference between this approach + //and a long-running RPC + stream, err := c.grpcClt.MapFn(ctx) + + // If any initial error, send it to the error channel + if err != nil { + go func(sErr error) { + errCh <- sdkerr.ToUDFErr("c.grpcClt.MapFn", sErr) + }(err) + } + + // Response routine: + // read the response from the server stream and send it to responseCh channel + // any error is sent to errCh channel + go func() { + // close this channel to indicate that no more elements left to receive from grpc + // We do defer here on the whole go-routine as even during a error scenario, we + // want to close the channel and stop forwarding any more responses from the UDF + // as we would be replaying the current ones. + defer close(responseCh) + for { + select { + // In case of the context done, return the error and stop further processing + case <-ctx.Done(): + errCh <- ctx.Err() + return + default: + var resp *flatmappb.MapResponse + resp, err := stream.Recv() + // check if this is EOF error, which indicates that no more responses left to process on the + // stream from the UDF, in such a case we return without any error to indicate this + if errors.Is(err, io.EOF) { + return + } + // If this is some other error, propagate it to error channel, + // also close the response channel(done using the defer close) to indicate no more messages being read + errSDK := sdkerr.ToUDFErr("flatmap c.grpcClt.MapFn", err) + if errSDK != nil { + errCh <- errSDK + return + } + //log.Print("MYDEBUG GOT FROM grpc ", resp.Result.GetUuid()) + responseCh <- resp + } + } + }() + // Read from the read messages and send them individually to the bi-di stream for processing + // in case there is an error in sending, send it to the error channel for handling + go func() { + for inputMsg := range datumStreamCh { + err := stream.Send(inputMsg) + if err != nil { + go func(sErr error) { + errCh <- sdkerr.ToUDFErr("c.grpcClt.MapFn", sErr) + }(err) + break + } + } + // CloseSend closes the send direction of the stream. This indicates to the + // UDF that we have sent all requests from the client, and it can safely + // stop listening on the stream + sendErr := stream.CloseSend() + if sendErr != nil && !errors.Is(sendErr, io.EOF) { + go func(sErr error) { + errCh <- sdkerr.ToUDFErr("flatmap c.grpcClt.MapFn stream.CloseSend()", sErr) + }(sendErr) + } + }() + return responseCh, errCh + + // FOR BYPASSING GRPC + // Read from the read messages and send them individually to the bi-di stream for processing + // in case there is an error in sending, send it to the error channel for handling + //go func() { + // defer close(responseCh) + // //var wg sync.WaitGroup + // //for i := 0; i < 1000; i++ { + // // wg.Add(1) + // // go func() { + // // go TestMap(ctx, datumStreamCh, responseCh, &wg) + // // }() + // //} + // //wg.Wait() + // for inputMsg := range datumStreamCh { + // mapResp := TestMap(ctx, inputMsg) + // for _, resp := range mapResp { + // responseCh <- resp + // } + // } + //}() + //return responseCh, errCh +} + +func (c client) CloseConn(ctx context.Context) error { + return c.conn.Close() +} + +func (c client) IsReady(ctx context.Context, in *emptypb.Empty) (bool, error) { + resp, err := c.grpcClt.IsReady(ctx, in) + if err != nil { + return false, err + } + return resp.GetReady(), nil +} + +// New creates a new client object. +func New(serverInfo *info.ServerInfo, inputOptions ...sdkclient.Option) (Client, error) { + var opts = sdkclient.DefaultOptions(sdkclient.FlatmapAddr) + + for _, inputOption := range inputOptions { + inputOption(opts) + } + + // Connect to the server + conn, err := grpcutil.ConnectToServer(opts.UdsSockAddr(), serverInfo, opts.MaxMessageSize()) + if err != nil { + return nil, err + } + + c := new(client) + c.conn = conn + c.grpcClt = flatmappb.NewFlatmapClient(conn) + return c, nil +} + +func TestMap(ctx context.Context, msg *flatmappb.MapRequest) []*flatmappb.MapResponse { + val := msg.GetValue() + var elements []*flatmappb.MapResponse + strs := strings.Split(string(val), ",") + for idx, x := range strs { + elements = append(elements, &flatmappb.MapResponse{ + Result: &flatmappb.MapResponse_Result{ + Keys: msg.GetKeys(), + Value: []byte(x), + Tags: msg.GetKeys(), + EOR: false, + Uuid: msg.GetUuid(), + Index: strconv.Itoa(idx), + Total: int32(len(strs)), + }, + }) + //element := &flatmappb.MapResponse{ + // Result: &flatmappb.MapResponse_Result{ + // Keys: msg.GetKeys(), + // Value: []byte(x), + // Tags: msg.GetKeys(), + // EOR: false, + // Uuid: msg.GetUuid(), + // Index: strconv.Itoa(idx), + // Total: int32(len(strs)), + // }, + //} + //responseChan <- element + } + if len(strs) == 0 { + // Append the EOR to indicate that the processing for the given request has completed + elements = append(elements, &flatmappb.MapResponse{ + Result: &flatmappb.MapResponse_Result{ + EOR: true, + Uuid: msg.GetUuid(), + Total: int32(len(strs)), + }, + }) + //element := &flatmappb.MapResponse{ + // Result: &flatmappb.MapResponse_Result{ + // EOR: true, + // Uuid: msg.GetUuid(), + // Total: int32(len(strs)), + // }, + //} + //responseChan <- element + } + return elements +} diff --git a/pkg/sdkclient/flatmapper/interface.go b/pkg/sdkclient/flatmapper/interface.go new file mode 100644 index 0000000000..b039303080 --- /dev/null +++ b/pkg/sdkclient/flatmapper/interface.go @@ -0,0 +1,14 @@ +package flatmapper + +import ( + flatmappb "github.com/numaproj/numaflow-go/pkg/apis/proto/flatmap/v1" + "golang.org/x/net/context" + "google.golang.org/protobuf/types/known/emptypb" +) + +// Client contains methods to call a gRPC client. +type Client interface { + CloseConn(ctx context.Context) error + IsReady(ctx context.Context, in *emptypb.Empty) (bool, error) + MapFn(ctx context.Context, datumStreamCh <-chan *flatmappb.MapRequest) (<-chan *flatmappb.MapResponse, <-chan error) +} diff --git a/pkg/sdkclient/mapper/client.go b/pkg/sdkclient/mapper/client.go index d22c852906..8759e67cfd 100644 --- a/pkg/sdkclient/mapper/client.go +++ b/pkg/sdkclient/mapper/client.go @@ -18,12 +18,14 @@ package mapper import ( "context" + "strings" "google.golang.org/grpc" "google.golang.org/protobuf/types/known/emptypb" mappb "github.com/numaproj/numaflow-go/pkg/apis/proto/map/v1" "github.com/numaproj/numaflow-go/pkg/info" + "github.com/numaproj/numaflow/pkg/sdkclient" sdkerror "github.com/numaproj/numaflow/pkg/sdkclient/error" grpcutil "github.com/numaproj/numaflow/pkg/sdkclient/grpc" @@ -82,9 +84,27 @@ func (c *client) IsReady(ctx context.Context, in *emptypb.Empty) (bool, error) { // MapFn applies a function to each datum element. func (c *client) MapFn(ctx context.Context, request *mappb.MapRequest) (*mappb.MapResponse, error) { mapResponse, err := c.grpcClt.MapFn(ctx, request) + //mapResponse, err := TestMap(ctx, request) err = sdkerror.ToUDFErr("c.grpcClt.MapFn", err) if err != nil { return nil, err } return mapResponse, nil } + +func TestMap(ctx context.Context, msg *mappb.MapRequest) (*mappb.MapResponse, error) { + val := msg.GetValue() + var elements []*mappb.MapResponse_Result + strs := strings.Split(string(val), ",") + for _, x := range strs { + elements = append(elements, &mappb.MapResponse_Result{ + Keys: msg.GetKeys(), + Value: []byte(x), + Tags: msg.GetKeys(), + }) + } + res := &mappb.MapResponse{ + Results: elements, + } + return res, nil +} diff --git a/pkg/sinks/blackhole/blackhole.go b/pkg/sinks/blackhole/blackhole.go index 3ef0332123..fd82181268 100644 --- a/pkg/sinks/blackhole/blackhole.go +++ b/pkg/sinks/blackhole/blackhole.go @@ -34,6 +34,11 @@ type Blackhole struct { logger *zap.SugaredLogger } +func (b *Blackhole) WriteNew(ctx context.Context, message isb.Message) (isb.Offset, error) { + //TODO implement me + panic("implement me") +} + // NewBlackhole returns a new Blackhole sink. func NewBlackhole(ctx context.Context, vertexInstance *dfv1.VertexInstance) (*Blackhole, error) { return &Blackhole{ diff --git a/pkg/sinks/kafka/kafka.go b/pkg/sinks/kafka/kafka.go index cfbd2d27ff..b445fa35be 100644 --- a/pkg/sinks/kafka/kafka.go +++ b/pkg/sinks/kafka/kafka.go @@ -42,6 +42,11 @@ type ToKafka struct { log *zap.SugaredLogger } +func (tk *ToKafka) WriteNew(ctx context.Context, message isb.Message) (isb.Offset, error) { + //TODO implement me + panic("implement me") +} + // NewToKafka returns ToKafka type. func NewToKafka(ctx context.Context, vertexInstance *dfv1.VertexInstance) (*ToKafka, error) { diff --git a/pkg/sinks/logger/log.go b/pkg/sinks/logger/log.go index 3d7e31f9ac..774b54e415 100644 --- a/pkg/sinks/logger/log.go +++ b/pkg/sinks/logger/log.go @@ -37,6 +37,11 @@ type ToLog struct { logger *zap.SugaredLogger } +func (t *ToLog) WriteNew(ctx context.Context, message isb.Message) (isb.Offset, error) { + //TODO implement me + panic("implement me") +} + // NewToLog returns ToLog type. func NewToLog(ctx context.Context, vertexInstance *dfv1.VertexInstance) (*ToLog, error) { return &ToLog{ diff --git a/pkg/sinks/udsink/sink.go b/pkg/sinks/udsink/sink.go index 7e67a831c4..6cacab45d9 100644 --- a/pkg/sinks/udsink/sink.go +++ b/pkg/sinks/udsink/sink.go @@ -36,6 +36,11 @@ type UserDefinedSink struct { udsink SinkApplier } +func (s *UserDefinedSink) WriteNew(ctx context.Context, message isb.Message) (isb.Offset, error) { + //TODO implement me + panic("implement me") +} + // NewUserDefinedSink returns genericSink type. func NewUserDefinedSink(ctx context.Context, vertexInstance *dfv1.VertexInstance, udsink SinkApplier) (*UserDefinedSink, error) { return &UserDefinedSink{ diff --git a/pkg/udf/forward/forward.go b/pkg/udf/forward/forward.go index 89b91349e9..388db28f22 100644 --- a/pkg/udf/forward/forward.go +++ b/pkg/udf/forward/forward.go @@ -255,7 +255,7 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { // we fetch the watermark for the partition from which we read the message. processorWM := isdf.wmFetcher.ComputeWatermark(readMessages[0].ReadOffset, isdf.fromBufferPartition.GetPartitionIdx()) - var writeOffsets map[string][][]isb.Offset + //var writeOffsets map[string][][]isb.Offset if !isdf.opts.enableMapUdfStream { // create space for writeMessages specific to each step as we could forward to all the steps too. var messageToStep = make(map[string][][]isb.Message) @@ -297,7 +297,7 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { // wait till the processing is done. this will not be an infinite wait because the map UDF processing will exit if // context.Done() is closed. wg.Wait() - isdf.opts.logger.Debugw("concurrent applyUDF completed", zap.Int("concurrency", isdf.opts.udfConcurrency), zap.Duration("took", time.Since(concurrentUDFProcessingStart))) + isdf.opts.logger.Debugw("MYDEBUG: concurrent no udf completed", zap.Int("concurrency", isdf.opts.udfConcurrency), zap.Duration("took", time.Since(concurrentUDFProcessingStart))) metrics.ConcurrentUDFProcessingTime.With(map[string]string{metrics.LabelVertex: isdf.vertexName, metrics.LabelPipeline: isdf.pipelineName, metrics.LabelVertexType: string(dfv1.VertexTypeMapUDF), metrics.LabelVertexReplicaIndex: strconv.Itoa(int(isdf.vertexReplica))}).Observe(float64(time.Since(concurrentUDFProcessingStart).Microseconds())) // map UDF processing is done. @@ -324,7 +324,7 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { } // forward the message to the edge buffer (could be multiple edges) - writeOffsets, err = isdf.writeToBuffers(ctx, messageToStep) + _, err = isdf.writeToBuffers(ctx, messageToStep) if err != nil { isdf.opts.logger.Errorw("failed to write to toBuffers", zap.Error(err)) isdf.fromBufferPartition.NoAck(ctx, readOffsets) @@ -332,7 +332,7 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { } isdf.opts.logger.Debugw("writeToBuffers completed") } else { - writeOffsets, err = isdf.streamMessage(ctx, dataMessages, processorWM) + _, err = isdf.streamMessage(ctx, dataMessages, processorWM) if err != nil { isdf.opts.logger.Errorw("failed to streamMessage", zap.Error(err)) // As there's no partial failure, non-ack all the readOffsets @@ -345,23 +345,23 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { // a watermark in this batch processing cycle. // it's used to determine which buffers should receive an idle watermark. // It is created as a slice because it tracks per partition activity info. - var activeWatermarkBuffers = make(map[string][]bool) - // forward the highest watermark to all the edges to avoid idle edge problem - // TODO: sort and get the highest value - for toVertexName, toVertexBufferOffsets := range writeOffsets { - activeWatermarkBuffers[toVertexName] = make([]bool, len(toVertexBufferOffsets)) - if publisher, ok := isdf.wmPublishers[toVertexName]; ok { - for index, offsets := range toVertexBufferOffsets { - if len(offsets) > 0 { - publisher.PublishWatermark(processorWM, offsets[len(offsets)-1], int32(index)) - activeWatermarkBuffers[toVertexName][index] = true - // reset because the toBuffer partition is no longer idling - isdf.idleManager.MarkActive(isdf.fromBufferPartition.GetPartitionIdx(), isdf.toBuffers[toVertexName][index].GetName()) - } - // This (len(offsets) == 0) happens at conditional forwarding, there's no data written to the buffer - } - } - } + //var activeWatermarkBuffers = make(map[string][]bool) + //// forward the highest watermark to all the edges to avoid idle edge problem + //// TODO: sort and get the highest value + //for toVertexName, toVertexBufferOffsets := range writeOffsets { + // activeWatermarkBuffers[toVertexName] = make([]bool, len(toVertexBufferOffsets)) + // if publisher, ok := isdf.wmPublishers[toVertexName]; ok { + // for index, offsets := range toVertexBufferOffsets { + // if len(offsets) > 0 { + // publisher.PublishWatermark(processorWM, offsets[len(offsets)-1], int32(index)) + // activeWatermarkBuffers[toVertexName][index] = true + // // reset because the toBuffer partition is no longer idling + // isdf.idleManager.MarkActive(isdf.fromBufferPartition.GetPartitionIdx(), isdf.toBuffers[toVertexName][index].GetName()) + // } + // // This (len(offsets) == 0) happens at conditional forwarding, there's no data written to the buffer + // } + // } + //} // - condition1 "len(dataMessages) > 0" : // Meaning, we do have some data messages, but we may not have written to all out buffers or its partitions. // It could be all data messages are dropped, or conditional forwarding to part of the out buffers. @@ -371,19 +371,19 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { // send idle watermark only if we have idle out buffers // Note: When the len(dataMessages) is 0, meaning all the readMessages are control messages, we choose not to do extra steps // This is because, if the idle continues, we will eventually handle the idle watermark when we read the next batch where the len(readMessages) will be zero - if len(dataMessages) > 0 { - for bufferName := range isdf.wmPublishers { - for index, activePartition := range activeWatermarkBuffers[bufferName] { - if !activePartition { - // use the watermark of the current read batch for the idle watermark - // same as read len==0 because there's no event published to the buffer - if p, ok := isdf.wmPublishers[bufferName]; ok { - idlehandler.PublishIdleWatermark(ctx, isdf.fromBufferPartition.GetPartitionIdx(), isdf.toBuffers[bufferName][index], p, isdf.idleManager, isdf.opts.logger, isdf.vertexName, isdf.pipelineName, dfv1.VertexTypeMapUDF, isdf.vertexReplica, processorWM) - } - } - } - } - } + //if len(dataMessages) > 0 { + // for bufferName := range isdf.wmPublishers { + // for index, activePartition := range activeWatermarkBuffers[bufferName] { + // if !activePartition { + // // use the watermark of the current read batch for the idle watermark + // // same as read len==0 because there's no event published to the buffer + // if p, ok := isdf.wmPublishers[bufferName]; ok { + // idlehandler.PublishIdleWatermark(ctx, isdf.fromBufferPartition.GetPartitionIdx(), isdf.toBuffers[bufferName][index], p, isdf.idleManager, isdf.opts.logger, isdf.vertexName, isdf.pipelineName, dfv1.VertexTypeMapUDF, isdf.vertexReplica, processorWM) + // } + // } + // } + // } + //} // when we apply udf, we don't handle partial errors (it's either non or all, non will return early), // so we should be able to ack all the readOffsets including data messages and control messages @@ -398,6 +398,8 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { // ProcessingTimes of the entire forwardAChunk metrics.ForwardAChunkProcessingTime.With(map[string]string{metrics.LabelVertex: isdf.vertexName, metrics.LabelPipeline: isdf.pipelineName, metrics.LabelVertexType: string(dfv1.VertexTypeMapUDF), metrics.LabelVertexReplicaIndex: strconv.Itoa(int(isdf.vertexReplica))}).Observe(float64(time.Since(start).Microseconds())) + isdf.opts.logger.Debugw("MYDEBUG: total with udf time", zap.Int("concurrency", isdf.opts.udfConcurrency), zap.Duration("took", time.Since(start))) + } // streamMessage streams the data messages to the next step. @@ -551,13 +553,13 @@ func (isdf *InterStepDataForward) writeToBuffers( ) (writeOffsets map[string][][]isb.Offset, err error) { // messageToStep contains all the to buffers, so the messages could be empty (conditional forwarding). // So writeOffsets also contains all the to buffers, but the returned offsets might be empty. - writeOffsets = make(map[string][][]isb.Offset) - for toVertexName, toVertexMessages := range messageToStep { - writeOffsets[toVertexName] = make([][]isb.Offset, len(toVertexMessages)) - } + //writeOffsets = make(map[string][][]isb.Offset) + //for toVertexName, toVertexMessages := range messageToStep { + // writeOffsets[toVertexName] = make([][]isb.Offset, len(toVertexMessages)) + //} for toVertexName, toVertexBuffer := range isdf.toBuffers { for index, partition := range toVertexBuffer { - writeOffsets[toVertexName][index], err = isdf.writeToBuffer(ctx, partition, messageToStep[toVertexName][index]) + _, err = isdf.writeToBuffer(ctx, partition, messageToStep[toVertexName][index]) if err != nil { return nil, err }