From 6193bdac3b0dab6f729ccd621827657d7a931cd2 Mon Sep 17 00:00:00 2001 From: Sidhant Kohli Date: Sun, 2 Jun 2024 14:05:54 -0700 Subject: [PATCH 01/14] skeleton Signed-off-by: Sidhant Kohli --- cmd/commands/processor.go | 13 + pkg/apis/numaflow/v1alpha1/const.go | 3 + pkg/apis/numaflow/v1alpha1/vertex_types.go | 9 + pkg/flatmap/common.go | 171 ++ pkg/flatmap/flatmap_udf.go | 279 ++++ pkg/flatmap/forward/applier/mapper.go | 37 + pkg/flatmap/forward/applier/mapstreamer.go | 36 + pkg/flatmap/forward/doc.go | 15 + pkg/flatmap/forward/forward.go | 867 ++++++++++ pkg/flatmap/forward/forward_test.go | 1741 ++++++++++++++++++++ pkg/flatmap/forward/options.go | 92 ++ pkg/flatmap/forward/shutdown.go | 74 + pkg/flatmap/forward/shutdown_test.go | 156 ++ pkg/flatmap/rpc/errors.go | 51 + pkg/flatmap/rpc/grpc_flatmap.go | 219 +++ pkg/flatmap/tracker/tracker.go | 96 ++ pkg/flatmap/types/types.go | 20 + pkg/sdkclient/const.go | 2 + pkg/sdkclient/flatmapper/client.go | 139 ++ pkg/sdkclient/flatmapper/interface.go | 14 + 20 files changed, 4034 insertions(+) create mode 100644 pkg/flatmap/common.go create mode 100644 pkg/flatmap/flatmap_udf.go create mode 100644 pkg/flatmap/forward/applier/mapper.go create mode 100644 pkg/flatmap/forward/applier/mapstreamer.go create mode 100644 pkg/flatmap/forward/doc.go create mode 100644 pkg/flatmap/forward/forward.go create mode 100644 pkg/flatmap/forward/forward_test.go create mode 100644 pkg/flatmap/forward/options.go create mode 100644 pkg/flatmap/forward/shutdown.go create mode 100644 pkg/flatmap/forward/shutdown_test.go create mode 100644 pkg/flatmap/rpc/errors.go create mode 100644 pkg/flatmap/rpc/grpc_flatmap.go create mode 100644 pkg/flatmap/tracker/tracker.go create mode 100644 pkg/flatmap/types/types.go create mode 100644 pkg/sdkclient/flatmapper/client.go create mode 100644 pkg/sdkclient/flatmapper/interface.go 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..4128ab98a5 --- /dev/null +++ b/pkg/flatmap/flatmap_udf.go @@ -0,0 +1,279 @@ +/* +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, int(*u.VertexInstance.Vertex.Spec.Limits.ReadBatchSize)) + + // 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 + } + + // Iterate through the edges + for _, edge := range u.VertexInstance.Vertex.Spec.ToEdges { + 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..61d4987d01 --- /dev/null +++ b/pkg/flatmap/forward/applier/mapper.go @@ -0,0 +1,37 @@ +/* +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/flatmap/types" + "github.com/numaproj/numaflow/pkg/isb" +) + +// FlatmapApplier applies the map UDF on the read message and gives back a new message. Any UserError will be retried here, while +// InternalErr can be returned and could be retried by the callee. +type FlatmapApplier interface { + ApplyMap(ctx context.Context, messageStream []*isb.ReadMessage, writeChan chan<- *types.ResponseFlatmap) <-chan error +} + +// ApplyFlatmapFunc utility function used to create a FlatmapApplier implementation +type ApplyFlatmapFunc func(ctx context.Context, messageStream []*isb.ReadMessage, writeChan chan<- *types.ResponseFlatmap) <-chan error + +func (f ApplyFlatmapFunc) ApplyMap(ctx context.Context, messageStream []*isb.ReadMessage, writeChan chan<- *types.ResponseFlatmap) <-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..75f6ee8906 --- /dev/null +++ b/pkg/flatmap/forward/forward.go @@ -0,0 +1,867 @@ +/* +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" + "sync" + "time" + + "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/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 + 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, + } + + // 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 +} + +// readWriteMessagePair represents a read message and its processed (via map UDF) write messages. +type readWriteMessagePair struct { + readMessage *isb.ReadMessage + writeMessages []*isb.WriteMessage + udfError error +} + +// forwardAChunk forwards a chunk of message 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 only we are successfully able to ack +// the message 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. +func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { + start := time.Now() + // 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 { + 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 + //// 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 + } + + var dataMessages = make([]*isb.ReadMessage, 0, len(readMessages)) + + // 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 { + dataMessages = append(dataMessages, m) + } else { + isdf.opts.logger.Info("MYDEBUG: NEW non data ", m.ReadOffset) + } + + } + 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))) + + // TODO(stream): check watermark here + //// 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. + processorWM := isdf.wmFetcher.ComputeWatermark(readMessages[0].ReadOffset, isdf.fromBufferPartition.GetPartitionIdx()) + + //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) + // for toVertex := range isdf.toBuffers { + // // over allocating to have a predictable pattern + // messageToStep[toVertex] = make([][]isb.Message, len(isdf.toBuffers[toVertex])) + // } + // + // // udf concurrent processing request channel + // udfCh := make(chan *readWriteMessagePair) + // // udfResults stores the results after map UDF processing for all read messages. It indexes + // // a read message to the corresponding write message + // udfResults := make([]readWriteMessagePair, len(dataMessages)) + // // applyUDF, if there is an Internal error it is a blocking call and will return only if shutdown has been initiated. + // + // // create a pool of map UDF Processors + // var wg sync.WaitGroup + // for i := 0; i < isdf.opts.udfConcurrency; i++ { + // wg.Add(1) + // go func() { + // defer wg.Done() + // isdf.concurrentApplyUDF(ctx, udfCh) + // }() + // } + // concurrentUDFProcessingStart := time.Now() + // + // // send to map UDF only the data messages + // for idx, m := range dataMessages { + // // emit message size metric + // metrics.ReadBytesCount.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(m.Payload))) + // // assign watermark to the message + // m.Watermark = time.Time(processorWM) + // // send map UDF processing work to the channel + // udfResults[idx].readMessage = m + // udfCh <- &udfResults[idx] + // } + // // let the go routines know that there is no more work + // close(udfCh) + // // 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))) + // 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. + // + // // let's figure out which vertex to send the results to. + // // update the toBuffer(s) with writeMessages. + // for _, m := range udfResults { + // // look for errors in udf processing, if we see even 1 error NoAck all messages + // // then return. Handling partial retrying is not worth ATM. + // if m.udfError != nil { + // metrics.UDFError.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() + // isdf.opts.logger.Errorw("failed to applyUDF", zap.Error(m.udfError)) + // // As there's no partial failure, non-ack all the readOffsets + // isdf.fromBufferPartition.NoAck(ctx, readOffsets) + // return + // } + // // update toBuffers + // for _, message := range m.writeMessages { + // if err := isdf.whereToStep(message, messageToStep, m.readMessage); err != nil { + // isdf.opts.logger.Errorw("failed in whereToStep", zap.Error(err)) + // isdf.fromBufferPartition.NoAck(ctx, readOffsets) + // return + // } + // } + // } + // + // // forward the message to the edge buffer (could be multiple edges) + // writeOffsets, 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) + // return + // } + // isdf.opts.logger.Debugw("writeToBuffers completed") + //} else { + + isdf.opts.logger.Info("Read the messages") + 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 + isdf.fromBufferPartition.NoAck(ctx, readOffsets) + return + } + //} + + // TODO(stream): WATERMARKING THIS + + //TODO: sort and get the highest value + //// activeWatermarkBuffers records the buffers that the publisher has published + //// 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 + // } + // } + //} + //// - 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. + //// If we don't have this condition check, when dataMessages is zero but ctrlMessages > 0, we will + //// wrongly publish an idle watermark without the ctrl message and the ctrl message tracking map. + //// - condition 2 "len(activeWatermarkBuffers) < len(isdf.wmPublishers)" : + //// 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) + // } + // } + // } + // } + //} + + // 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 + //err = isdf.ackFromBuffer(ctx, readOffsets) + //// implicit return for posterity :-) + //if err != nil { + // isdf.opts.logger.Errorw("Failed to ack from buffer", zap.Error(err)) + // 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()}).Add(float64(len(readOffsets))) + // return + //} + //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(readOffsets))) + + // 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("concurrent applyUDF completed", zap.Int("concurrency", isdf.opts.udfConcurrency), zap.Duration("took", time.Since(start))) +} + +// streamMessage streams the data messages to the next step. +func (isdf *InterStepDataForward) streamMessage( + ctx context.Context, + dataMessages []*isb.ReadMessage, + processorWM wmb.Watermark, +) error { + // create space for writeMessages specific to each step as we could forward to all the steps too. + // these messages are for per partition (due to round-robin writes) for load balancing + 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])) + } + + if len(dataMessages) >= 1 { + + // emit message size metric + metrics.ReadBytesCount.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[0].Payload))) + // assign watermark to the message + dataMessages[0].Watermark = time.Time(processorWM) + + // process the mapStreamUDF and get the result + //start := time.Now() + metrics.UDFReadMessagesCount.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() + + writeMessageCh, errCh := isdf.invokeUDF(ctx, dataMessages) + + ackChan := isdf.invokeWriter(ctx, writeMessageCh, errCh) + + // TODO(stream): WATERMARKING THIS WRITES + + done := isdf.invokeAck(ctx, ackChan) + + <-done + + // TODO(stream): LETS START THIS WRITER IN A GOROUTINE + // Stream the message to the next vertex. First figure out which vertex + // to send the result to. Then update the toBuffer(s) with writeMessage. + //msgIndex := 0 + //for writeMessage := range writeMessageCh { + // writeMessage.Headers = dataMessages[0].Headers + // // add vertex name to the ID, since multiple vertices can publish to the same vertex and we need uniqueness across them + // writeMessage.ID = fmt.Sprintf("%s-%s-%d", dataMessages[0].ReadOffset.String(), isdf.vertexName, msgIndex) + // msgIndex += 1 + // 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(1)) + // + // // update toBuffers + // if err := isdf.whereToStep(&writeMessage, messageToStep, dataMessages[0]); err != nil { + // return nil, fmt.Errorf("failed at whereToStep, error: %w", err) + // } + // + // // Forward the message to the edge buffer (could be multiple edges) + // curWriteOffsets, err := isdf.writeToBuffers(ctx, messageToStep) + // if err != nil { + // return nil, fmt.Errorf("failed to write to toBuffers, error: %w", err) + // } + // // Merge curWriteOffsets into writeOffsets + // for vertexName, toVertexBufferOffsets := range curWriteOffsets { + // for index, offsets := range toVertexBufferOffsets { + // writeOffsets[vertexName][index] = append(writeOffsets[vertexName][index], offsets...) + // } + // } + //} + + // look for errors in udf processing, if we see even 1 error NoAck all messages + // then return. Handling partial retrying is not worth ATM. + //if err := errs.Wait(); err != nil { + // metrics.UDFError.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() + // // We do not retry as we are streaming + // if ok, _ := isdf.IsShuttingDown(); ok { + // isdf.opts.logger.Errorw("flatmapUDF.Apply, Stop called while stuck on an internal error", zap.Error(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 fmt.Errorf("failed to applyUDF, error: %w", err) + //} + + //metrics.UDFProcessingTime.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())) + } else { + // Even not data messages, forward the message to the edge buffer (could be multiple edges) + var err error + _, err = isdf.writeToBuffers(ctx, messageToStep) + if err != nil { + return fmt.Errorf("failed to write to toBuffers, error: %w", err) + } + } + + return nil +} + +// invokeUDF reads requests from the supplied PBQ, invokes the UDF to gets the response and writes the response to the +// main channel. +func (isdf *InterStepDataForward) invokeUDF(ctx context.Context, dataMessages []*isb.ReadMessage) (<-chan *types.ResponseFlatmap, <-chan error) { + writeChan := make(chan *types.ResponseFlatmap) + errCh := isdf.flatmapUDF.ApplyMap(ctx, dataMessages, writeChan) + return writeChan, errCh + // + //outerLoop: + // for { + // select { + // case err := <-errCh: + // if errors.Is(err, context.Canceled) || ctx.Err() != nil { + // isdf.opts.logger.Infow("Context is canceled, stopping the processAndForward", zap.Error(err)) + // return nil, nil + // } + // if err != nil { + // isdf.opts.logger.Error("Got an error while invoking ApplyReduce", zap.Error(err)) + // // TODO(stream): trigger shutdown + // } + // } + // } +} + +func (isdf *InterStepDataForward) ackRoutine(ctx context.Context, ackMsgChan <-chan *isb.ReadMessage, wg *sync.WaitGroup) { + defer wg.Done() +forwardLoop: + for { + select { + case <-ctx.Done(): + break forwardLoop + + case response, ok := <-ackMsgChan: + if !ok { + break forwardLoop + } + ackMessages := []isb.Offset{response.ReadOffset} + if err := isdf.ackFromBuffer(ctx, ackMessages); err != nil { + isdf.opts.logger.Error("MYDEBUG: ERROR IN ACK ", zap.Error(err)) + // TODO(stream): we have retried in the ackFromBuffer, should we trigger + // shutdown here then? + } + //isdf.opts.logger.Info("MYDEBUG: NEW GOT TO ACK THIS ", response.ReadOffset) + } + } +} +func (isdf *InterStepDataForward) invokeAck(ctx context.Context, ackMsgChan <-chan *isb.ReadMessage) (doneChan chan struct{}) { + logger := isdf.opts.logger + logger.Info("MYDEBUG: NO WG ACK ROUTINE ", isdf.opts.readBatchSize) + 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 +} + +func (isdf *InterStepDataForward) writeRoutine(ctx context.Context, writeMessageCh <-chan *types.ResponseFlatmap, errCh <-chan error, ackChan chan<- *isb.ReadMessage, wg *sync.WaitGroup) { + defer wg.Done() +outerLoop: + for { + select { + case err := <-errCh: + if errors.Is(err, context.Canceled) || ctx.Err() != nil { + isdf.opts.logger.Infow("Context is canceled, stopping the writer", zap.Error(err)) + return + } + if err != nil { + isdf.opts.logger.Error("MYDEBUG: NEW ERROR IN ERROR CH WRITE", zap.Error(err)) + // TODO(stream): got a error from the UDF, time to start the shutdown then or retry + //isdf.opts.logger.Panic("Got an error while invoking ApplyMap", zap.Error(err)) + } + case response, ok := <-writeMessageCh: + if !ok { + break outerLoop + } + //isdf.opts.logger.Info("MYDEBUG: NEW SENDING TO ACK ", response.ParentMessage.ReadOffset, " ", response.Uid) + 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 := response.RespMessage + writeMessages := []*isb.WriteMessage{writeMessage} + if err := isdf.forwardToBuffers(ctx, &writeMessages, response.ParentMessage, messageToStep); err != nil { + // // TODO(stream): mark as no ack directly or retry + isdf.opts.logger.Error("MYDEBUG: NEW ERROR IN WRITE", zap.Error(err)) + } + + 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(1)) + ackChan <- response.ParentMessage + } + } + +} + +func (isdf *InterStepDataForward) invokeWriter(ctx context.Context, writeMessageCh <-chan *types.ResponseFlatmap, errCh <-chan error) <-chan *isb.ReadMessage { + ackChan := make(chan *isb.ReadMessage) + go func() { + defer close(ackChan) + group := sync.WaitGroup{} + for i := 0; i < int(isdf.opts.readBatchSize); i++ { + group.Add(1) + go isdf.writeRoutine(ctx, writeMessageCh, errCh, ackChan, &group) + } + group.Wait() + }() + return ackChan +} + +func (isdf *InterStepDataForward) forwardToBuffers(ctx context.Context, writeMessages *[]*isb.WriteMessage, readMessage *isb.ReadMessage, messageToStep map[string][]isb.Message) error { + if len(*writeMessages) == 0 { + return nil + } + for _, message := range *writeMessages { + if err := isdf.whereToStep(message, 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 +} + +// ackFromBuffer acknowledges an array of offsets back to fromBufferPartition and is a blocking call or until shutdown has been initiated. +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 { + writeOffsets[toVertexName][index], err = isdf.writeToBuffer(ctx, partition, messageToStep[toVertexName][index]) + if err != nil { + return nil, err + } + } + } + return writeOffsets, 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.Write(ctx, []isb.Message{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 = errs[0]; 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", errorArrayToMap(errs)), + 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 +} + +//// concurrentApplyUDF applies the map UDF based on the request from the channel +//func (isdf *InterStepDataForward) concurrentApplyUDF(ctx context.Context, readMessagePair <-chan *readWriteMessagePair) { +// for message := range readMessagePair { +// start := time.Now() +// metrics.UDFReadMessagesCount.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() +// writeMessages, err := isdf.applyUDF(ctx, message.readMessage) +// 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(len(writeMessages))) +// // set the headers for the write messages +// for _, m := range writeMessages { +// m.Headers = message.readMessage.Headers +// } +// message.writeMessages = append(message.writeMessages, writeMessages...) +// message.udfError = err +// metrics.UDFProcessingTime.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())) +// } +//} + +//// applyUDF applies the map UDF and will block if there is any InternalErr. On the other hand, if this is a UserError +//// the skip flag is set. ShutDown flag will only if there is an InternalErr and ForceStop has been invoked. +//// The UserError retry will be done on the ApplyUDF. +//func (isdf *InterStepDataForward) applyUDF(ctx context.Context, readMessage *isb.ReadMessage) ([]*isb.WriteMessage, error) { +// for { +// writeMessages, err := isdf.flatmapUDF.ApplyMap(ctx, readMessage) +// if err != nil { +// isdf.opts.logger.Errorw("flatmapUDF.Apply error", zap.Error(err)) +// // TODO: implement retry with backoff etc. +// time.Sleep(isdf.opts.retryInterval) +// // keep retrying, I cannot think of a use case where a user could say, errors are fine :-) +// // as a platform we should not lose or corrupt data. +// // this does not mean we should prohibit this from a shutdown. +// if ok, _ := isdf.IsShuttingDown(); ok { +// isdf.opts.logger.Errorw("flatmapUDF.Apply, Stop called while stuck on an internal error", zap.Error(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 nil, err +// } +// continue +// } else { +// for index, m := range writeMessages { +// // add vertex name to the ID, since multiple vertices can publish to the same vertex and we need uniqueness across them +// m.ID = fmt.Sprintf("%s-%s-%d", readMessage.ReadOffset.String(), isdf.vertexName, index) +// } +// return writeMessages, 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 +} 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..fa0fe248f2 --- /dev/null +++ b/pkg/flatmap/rpc/grpc_flatmap.go @@ -0,0 +1,219 @@ +package rpc + +import ( + "fmt" + "log" + "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" + "github.com/numaproj/numaflow/pkg/isb" + 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 + readBatchSize int + idx int +} + +func NewUDSgRPCBasedFlatmap(client flatmapper.Client, batchSize int) *GRPCBasedFlatmap { + return &GRPCBasedFlatmap{client: client, tracker: tracker.NewTracker(), readBatchSize: batchSize} +} + +// 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) + } + } + } +} + +func (u *GRPCBasedFlatmap) ApplyMap(ctx context.Context, messageStream []*isb.ReadMessage, responseCh chan<- *types.ResponseFlatmap) <-chan error { + var ( + errCh = make(chan error) + ) + flatmapRequests := make(chan *flatmappb.MapRequest) + defer close(flatmapRequests) + + // invoke the AsyncReduceFn method with requestCh channel and send the result to responseCh channel + // and any error to errCh channel + go func() { + // TODO(stream): check this close here + defer close(responseCh) + resultCh, reduceErrCh := u.client.MapFn(ctx, flatmapRequests) + for { + select { + case result, ok := <-resultCh: + if !ok || result == nil { + //// if the resultCh channel is closed, close the responseCh + return + } + resp, remove := u.parseMapResponse(result) + if remove { + u.tracker.RemoveRequest(resp.Uid) + continue + } + responseCh <- resp + log.Println("MYDEBUG: sending to writeCH", resp.Uid) + case err := <-reduceErrCh: + // ctx.Done() event will be handled by the AsyncReduceFn method + // so we don't need a separate case for ctx.Done() here + if err == ctx.Err() { + errCh <- err + return + } + if err != nil { + errCh <- convertToUdfError(err) + // TODO(stream): check error here + } + } + } + }() + + for _, req := range messageStream { + d := u.createMapRequest(req) + log.Print("MYDEBUG: TRYING TO SEND HERE", d.Uuid) + flatmapRequests <- d + } + + return errCh +} + +func (u *GRPCBasedFlatmap) createMapRequest(msg *isb.ReadMessage) *flatmappb.MapRequest { + keys := msg.Keys + payload := msg.Body.Payload + + uid := u.tracker.AddRequest(msg) + + var d = &flatmappb.MapRequest{ + Keys: keys, + Value: payload, + EventTime: timestamppb.New(msg.EventTime), + Watermark: timestamppb.New(msg.Watermark), + Headers: msg.Headers, + Uuid: uid, + } + return d +} +func (u *GRPCBasedFlatmap) parseMapResponse(resp *flatmappb.MapResponse) (parsedResp *types.ResponseFlatmap, requestDone bool) { + result := resp.Result + eor := result.GetEOR() + uid := result.GetUuid() + parentRequest, ok := u.tracker.GetRequest(uid) + // TODO(stream): check what should be path for !ok + if !ok { + // u.tracker.NewResponse(uid) + // idx = 1 + } + // Request has completed remove from the tracker module + if eor == true { + return nil, true + } + //idx, present := u.tracker.GetIdx(uid) + //if !present { + + //} + keys := result.GetKeys() + taggedMessage := &isb.WriteMessage{ + Message: isb.Message{ + Header: isb.Header{ + MessageInfo: parentRequest.MessageInfo, + // TODO(stream): Check what will be the unique ID to use here + //msgId := fmt.Sprintf("%s-%d-%s-%d", u.vertexName, u.vertexReplica, partitionID.String(), index) + //ID: fmt.Sprintf("%s-%s-%d", dataMessages[0].ReadOffset.String(), isdf.vertexName, msgIndex) + ID: fmt.Sprintf("%s-%d", parentRequest.ReadOffset.String(), u.idx), + Keys: keys, + }, + Body: isb.Body{ + Payload: result.GetValue(), + }, + }, + + Tags: result.GetTags(), + } + + u.idx += 1 + //u.tracker.IncrementRespIdx(uid) + return &types.ResponseFlatmap{ + ParentMessage: parentRequest, + Uid: uid, + RespMessage: taggedMessage, + }, false +} + +// convertToUdfError converts the error returned by the reduceFn to ApplyUDFErr +func convertToUdfError(err error) error { + // if any error happens in reduce + // will exit and restart the numa container + udfErr, _ := sdkerr.FromError(err) + switch udfErr.ErrorKind() { + case sdkerr.Retryable: + // TODO: currently we don't handle retryable errors for reduce + return &ApplyUDFErr{ + UserUDFErr: false, + Message: fmt.Sprintf("gRPC client.ReduceFn failed, %s", err), + InternalErr: InternalErr{ + Flag: true, + MainCarDown: false, + }, + } + case sdkerr.NonRetryable: + return &ApplyUDFErr{ + UserUDFErr: false, + Message: fmt.Sprintf("gRPC client.ReduceFn 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.ReduceFn 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..29a73d4c4b --- /dev/null +++ b/pkg/flatmap/tracker/tracker.go @@ -0,0 +1,96 @@ +package tracker + +import ( + "log" + "sync" + + "github.com/google/uuid" + + "github.com/numaproj/numaflow/pkg/isb" +) + +type Tracker struct { + requestMap sync.Map + // TODO(stream): check if this will be inefficient, and there is a better way for this + responseIdx sync.Map + lock sync.RWMutex + m map[string]*isb.ReadMessage +} + +func NewTracker() *Tracker { + return &Tracker{ + requestMap: sync.Map{}, + responseIdx: sync.Map{}, + m: make(map[string]*isb.ReadMessage), + } +} + +func GetNewId() string { + id, _ := uuid.NewUUID() + return id.String() +} + +func (t *Tracker) AddRequest(msg *isb.ReadMessage) string { + id := GetNewId() + //t.requestMap.Store(id, msg) + t.Set(id, msg) + return id +} + +func (t *Tracker) GetRequest(id string) (*isb.ReadMessage, 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() + for k, v := range t.m { + log.Println("MYDEBUG: MAP VALS", k, " ", v) + } +} + +func (t *Tracker) Get(key string) (*isb.ReadMessage, bool) { + t.lock.RLock() + defer t.lock.RUnlock() + item, ok := t.m[key] + return item, ok +} + +func (t *Tracker) Set(key string, value *isb.ReadMessage) { + 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) +} diff --git a/pkg/flatmap/types/types.go b/pkg/flatmap/types/types.go new file mode 100644 index 0000000000..e0a08a1228 --- /dev/null +++ b/pkg/flatmap/types/types.go @@ -0,0 +1,20 @@ +package types + +import "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? +type ResponseFlatmap struct { + ParentMessage *isb.ReadMessage + Uid string + RespMessage *isb.WriteMessage +} + +type WriteMsgFlatmap struct { + Message *ResponseFlatmap + AckIt bool +} + +type AckMsgFlatmap struct { + Message *isb.ReadMessage + AckIt bool +} 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..d0c556a3b2 --- /dev/null +++ b/pkg/sdkclient/flatmapper/client.go @@ -0,0 +1,139 @@ +/* +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" + "log" + + 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 +} + +func (c client) MapFn(ctx context.Context, datumStreamCh <-chan *flatmappb.MapRequest) (<-chan *flatmappb.MapResponse, <-chan error) { + var ( + errCh = make(chan error) + responseCh = make(chan *flatmappb.MapResponse) + ) + + log.Println("MYDEBUG: WHATS HAPPENING", len(datumStreamCh), datumStreamCh) + + // stream the messages to server + stream, err := c.grpcClt.MapFn(ctx) + + // read the response from the server stream and send it to responseCh channel + // any error is sent to errCh channel + go func() { + defer close(responseCh) + for { + select { + case <-ctx.Done(): + errCh <- ctx.Err() + return + default: + var resp *flatmappb.MapResponse + resp, err = stream.Recv() + //if err == io.EOF { + // return + //} + if errors.Is(err, io.EOF) { + // skip selection on nil channel + //errCh = nil + //close(responseCh) + return + } + errSDK := sdkerr.ToUDFErr("c.grpcClt.MapStreamFn", err) + if errSDK != nil { + log.Println("MYDEBUG: ERROR in recv", err, errSDK) + errCh <- errSDK + return + } + log.Println("MYDEBUG: GOT IT FROM GRPC", resp.Result.Uuid) + responseCh <- resp + } + } + }() + // Read from the inputStream and send messages + for inputMsg := range datumStreamCh { + log.Println("MYDEBUG: Sending to grpc", inputMsg.Uuid) + err = stream.Send(inputMsg) + if err != nil { + go func(sErr error) { + errCh <- sdkerr.ToUDFErr("c.grpcClt.MapFn", sErr) + }(err) + break + } + //if err != nil { + // errCh <- sdkerr.ToUDFErr("MapFn stream.Send()", sendErr) + // return + //} + } + sendErr := stream.CloseSend() + if sendErr != nil && !errors.Is(sendErr, io.EOF) { + go func(sErr error) { + errCh <- sdkerr.ToUDFErr("c.grpcClt.MapFn stream.CloseSend()", sErr) + }(sendErr) + } + + 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 +} 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) +} From 8a7449c73829ee3f4d9bccac9400e1bd0c46af3a Mon Sep 17 00:00:00 2001 From: Sidhant Kohli Date: Sun, 2 Jun 2024 15:48:25 -0700 Subject: [PATCH 02/14] init Signed-off-by: Sidhant Kohli --- pkg/flatmap/forward/forward.go | 160 +++----------------------------- pkg/flatmap/rpc/grpc_flatmap.go | 2 +- 2 files changed, 12 insertions(+), 150 deletions(-) diff --git a/pkg/flatmap/forward/forward.go b/pkg/flatmap/forward/forward.go index 75f6ee8906..a469a53c7c 100644 --- a/pkg/flatmap/forward/forward.go +++ b/pkg/flatmap/forward/forward.go @@ -247,142 +247,19 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { //// 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. - processorWM := isdf.wmFetcher.ComputeWatermark(readMessages[0].ReadOffset, isdf.fromBufferPartition.GetPartitionIdx()) - - //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) - // for toVertex := range isdf.toBuffers { - // // over allocating to have a predictable pattern - // messageToStep[toVertex] = make([][]isb.Message, len(isdf.toBuffers[toVertex])) - // } - // - // // udf concurrent processing request channel - // udfCh := make(chan *readWriteMessagePair) - // // udfResults stores the results after map UDF processing for all read messages. It indexes - // // a read message to the corresponding write message - // udfResults := make([]readWriteMessagePair, len(dataMessages)) - // // applyUDF, if there is an Internal error it is a blocking call and will return only if shutdown has been initiated. - // - // // create a pool of map UDF Processors - // var wg sync.WaitGroup - // for i := 0; i < isdf.opts.udfConcurrency; i++ { - // wg.Add(1) - // go func() { - // defer wg.Done() - // isdf.concurrentApplyUDF(ctx, udfCh) - // }() - // } - // concurrentUDFProcessingStart := time.Now() - // - // // send to map UDF only the data messages - // for idx, m := range dataMessages { - // // emit message size metric - // metrics.ReadBytesCount.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(m.Payload))) - // // assign watermark to the message - // m.Watermark = time.Time(processorWM) - // // send map UDF processing work to the channel - // udfResults[idx].readMessage = m - // udfCh <- &udfResults[idx] - // } - // // let the go routines know that there is no more work - // close(udfCh) - // // 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))) - // 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. - // - // // let's figure out which vertex to send the results to. - // // update the toBuffer(s) with writeMessages. - // for _, m := range udfResults { - // // look for errors in udf processing, if we see even 1 error NoAck all messages - // // then return. Handling partial retrying is not worth ATM. - // if m.udfError != nil { - // metrics.UDFError.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() - // isdf.opts.logger.Errorw("failed to applyUDF", zap.Error(m.udfError)) - // // As there's no partial failure, non-ack all the readOffsets - // isdf.fromBufferPartition.NoAck(ctx, readOffsets) - // return - // } - // // update toBuffers - // for _, message := range m.writeMessages { - // if err := isdf.whereToStep(message, messageToStep, m.readMessage); err != nil { - // isdf.opts.logger.Errorw("failed in whereToStep", zap.Error(err)) - // isdf.fromBufferPartition.NoAck(ctx, readOffsets) - // return - // } - // } - // } - // - // // forward the message to the edge buffer (could be multiple edges) - // writeOffsets, 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) - // return - // } - // isdf.opts.logger.Debugw("writeToBuffers completed") - //} else { - isdf.opts.logger.Info("Read the messages") - 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 - isdf.fromBufferPartition.NoAck(ctx, readOffsets) - return - } - //} + //processorWM := isdf.wmFetcher.ComputeWatermark(readMessages[0].ReadOffset, isdf.fromBufferPartition.GetPartitionIdx()) - // TODO(stream): WATERMARKING THIS - - //TODO: sort and get the highest value - //// activeWatermarkBuffers records the buffers that the publisher has published - //// 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 - // } - // } - //} - //// - 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. - //// If we don't have this condition check, when dataMessages is zero but ctrlMessages > 0, we will - //// wrongly publish an idle watermark without the ctrl message and the ctrl message tracking map. - //// - condition 2 "len(activeWatermarkBuffers) < len(isdf.wmPublishers)" : - //// 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) - // } - // } - // } - // } - //} + // process UDF + writeMessageCh, udfErrorCh := isdf.invokeUDF(ctx, dataMessages) + + ackChan := isdf.invokeWriter(ctx, writeMessageCh, udfErrorCh) + + // TODO(stream): WATERMARKING THIS WRITES + + done := isdf.invokeAck(ctx, ackChan) + + <-done // 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 @@ -498,21 +375,6 @@ func (isdf *InterStepDataForward) invokeUDF(ctx context.Context, dataMessages [] writeChan := make(chan *types.ResponseFlatmap) errCh := isdf.flatmapUDF.ApplyMap(ctx, dataMessages, writeChan) return writeChan, errCh - // - //outerLoop: - // for { - // select { - // case err := <-errCh: - // if errors.Is(err, context.Canceled) || ctx.Err() != nil { - // isdf.opts.logger.Infow("Context is canceled, stopping the processAndForward", zap.Error(err)) - // return nil, nil - // } - // if err != nil { - // isdf.opts.logger.Error("Got an error while invoking ApplyReduce", zap.Error(err)) - // // TODO(stream): trigger shutdown - // } - // } - // } } func (isdf *InterStepDataForward) ackRoutine(ctx context.Context, ackMsgChan <-chan *isb.ReadMessage, wg *sync.WaitGroup) { diff --git a/pkg/flatmap/rpc/grpc_flatmap.go b/pkg/flatmap/rpc/grpc_flatmap.go index fa0fe248f2..3814bba83b 100644 --- a/pkg/flatmap/rpc/grpc_flatmap.go +++ b/pkg/flatmap/rpc/grpc_flatmap.go @@ -85,7 +85,7 @@ func (u *GRPCBasedFlatmap) ApplyMap(ctx context.Context, messageStream []*isb.Re continue } responseCh <- resp - log.Println("MYDEBUG: sending to writeCH", resp.Uid) + log.Println("MYDEBUG: sending to writeCh", resp.Uid) case err := <-reduceErrCh: // ctx.Done() event will be handled by the AsyncReduceFn method // so we don't need a separate case for ctx.Done() here From 3d03774f62ce519ec252e4116c5f3670cc7627e3 Mon Sep 17 00:00:00 2001 From: Sidhant Kohli Date: Mon, 3 Jun 2024 16:21:58 -0700 Subject: [PATCH 03/14] comments Signed-off-by: Sidhant Kohli --- pkg/flatmap/forward/applier/mapper.go | 6 +- pkg/flatmap/forward/forward.go | 211 ++++++-------------------- pkg/flatmap/rpc/grpc_flatmap.go | 125 +++++++++------ pkg/flatmap/tracker/tracker.go | 1 + pkg/flatmap/types/types.go | 2 + pkg/sdkclient/flatmapper/client.go | 76 ++++++---- 6 files changed, 183 insertions(+), 238 deletions(-) diff --git a/pkg/flatmap/forward/applier/mapper.go b/pkg/flatmap/forward/applier/mapper.go index 61d4987d01..cfb6ee2f10 100644 --- a/pkg/flatmap/forward/applier/mapper.go +++ b/pkg/flatmap/forward/applier/mapper.go @@ -23,9 +23,11 @@ import ( "github.com/numaproj/numaflow/pkg/isb" ) -// FlatmapApplier applies the map UDF on the read message and gives back a new message. Any UserError will be retried here, while -// InternalErr can be returned and could be retried by the callee. +// 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 []*isb.ReadMessage, writeChan chan<- *types.ResponseFlatmap) <-chan error } diff --git a/pkg/flatmap/forward/forward.go b/pkg/flatmap/forward/forward.go index a469a53c7c..3df9629f3b 100644 --- a/pkg/flatmap/forward/forward.go +++ b/pkg/flatmap/forward/forward.go @@ -174,23 +174,34 @@ type readWriteMessagePair struct { udfError error } -// forwardAChunk forwards a chunk of message 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 only we are successfully able to ack -// the message 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. +// 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 + // 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 done //// 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 { @@ -221,41 +232,58 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { // } // return //} + // TODO(stream): remove once idle watermark is done if len(readMessages) == 0 { return } + // 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([]*isb.ReadMessage, 0, len(readMessages)) - - // store the offsets of the messages we read from ISB + // 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 { dataMessages = append(dataMessages, m) - } else { - isdf.opts.logger.Info("MYDEBUG: NEW non data ", 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))) - // TODO(stream): check watermark here //// 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()) - // process UDF - writeMessageCh, udfErrorCh := isdf.invokeUDF(ctx, dataMessages) + // 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. + // writeChan 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 done, to indicate that no further processing is required. + // TODO(stream): on an error, we should close the writeChan, so that the writer can stop processing. + // 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 + writeChan := make(chan *types.ResponseFlatmap) - ackChan := isdf.invokeWriter(ctx, writeMessageCh, udfErrorCh) + // send the input messages for processing + // The error channel is used to signal any errors that might have occurred during the UDF processing. + // TODO(stream): add the error handler for graceful handling + udfErrorCh := isdf.flatmapUDF.ApplyMap(ctx, dataMessages, writeChan) - // TODO(stream): WATERMARKING THIS WRITES + // TODO(stream): Publish the watermark for these writers + ackChan := isdf.invokeWriter(ctx, writeChan, udfErrorCh) done := isdf.invokeAck(ctx, ackChan) @@ -277,105 +305,12 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { isdf.opts.logger.Debugw("concurrent applyUDF completed", zap.Int("concurrency", isdf.opts.udfConcurrency), zap.Duration("took", time.Since(start))) } -// streamMessage streams the data messages to the next step. -func (isdf *InterStepDataForward) streamMessage( - ctx context.Context, - dataMessages []*isb.ReadMessage, - processorWM wmb.Watermark, -) error { - // create space for writeMessages specific to each step as we could forward to all the steps too. - // these messages are for per partition (due to round-robin writes) for load balancing - 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])) - } - - if len(dataMessages) >= 1 { - - // emit message size metric - metrics.ReadBytesCount.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[0].Payload))) - // assign watermark to the message - dataMessages[0].Watermark = time.Time(processorWM) - - // process the mapStreamUDF and get the result - //start := time.Now() - metrics.UDFReadMessagesCount.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() - - writeMessageCh, errCh := isdf.invokeUDF(ctx, dataMessages) - - ackChan := isdf.invokeWriter(ctx, writeMessageCh, errCh) - - // TODO(stream): WATERMARKING THIS WRITES - - done := isdf.invokeAck(ctx, ackChan) - - <-done - - // TODO(stream): LETS START THIS WRITER IN A GOROUTINE - // Stream the message to the next vertex. First figure out which vertex - // to send the result to. Then update the toBuffer(s) with writeMessage. - //msgIndex := 0 - //for writeMessage := range writeMessageCh { - // writeMessage.Headers = dataMessages[0].Headers - // // add vertex name to the ID, since multiple vertices can publish to the same vertex and we need uniqueness across them - // writeMessage.ID = fmt.Sprintf("%s-%s-%d", dataMessages[0].ReadOffset.String(), isdf.vertexName, msgIndex) - // msgIndex += 1 - // 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(1)) - // - // // update toBuffers - // if err := isdf.whereToStep(&writeMessage, messageToStep, dataMessages[0]); err != nil { - // return nil, fmt.Errorf("failed at whereToStep, error: %w", err) - // } - // - // // Forward the message to the edge buffer (could be multiple edges) - // curWriteOffsets, err := isdf.writeToBuffers(ctx, messageToStep) - // if err != nil { - // return nil, fmt.Errorf("failed to write to toBuffers, error: %w", err) - // } - // // Merge curWriteOffsets into writeOffsets - // for vertexName, toVertexBufferOffsets := range curWriteOffsets { - // for index, offsets := range toVertexBufferOffsets { - // writeOffsets[vertexName][index] = append(writeOffsets[vertexName][index], offsets...) - // } - // } - //} - - // look for errors in udf processing, if we see even 1 error NoAck all messages - // then return. Handling partial retrying is not worth ATM. - //if err := errs.Wait(); err != nil { - // metrics.UDFError.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() - // // We do not retry as we are streaming - // if ok, _ := isdf.IsShuttingDown(); ok { - // isdf.opts.logger.Errorw("flatmapUDF.Apply, Stop called while stuck on an internal error", zap.Error(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 fmt.Errorf("failed to applyUDF, error: %w", err) - //} - - //metrics.UDFProcessingTime.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())) - } else { - // Even not data messages, forward the message to the edge buffer (could be multiple edges) - var err error - _, err = isdf.writeToBuffers(ctx, messageToStep) - if err != nil { - return fmt.Errorf("failed to write to toBuffers, error: %w", err) - } - } - - return nil -} - -// invokeUDF reads requests from the supplied PBQ, invokes the UDF to gets the response and writes the response to the -// main channel. -func (isdf *InterStepDataForward) invokeUDF(ctx context.Context, dataMessages []*isb.ReadMessage) (<-chan *types.ResponseFlatmap, <-chan error) { - writeChan := make(chan *types.ResponseFlatmap) - errCh := isdf.flatmapUDF.ApplyMap(ctx, dataMessages, writeChan) - return writeChan, errCh -} +// +//// invokeUDF reads requests from the supplied PBQ, invokes the UDF to gets the response and writes the response to the +//// main channel. +//func (isdf *InterStepDataForward) invokeUDF(ctx context.Context, dataMessages []*isb.ReadMessage) (<-chan *types.ResponseFlatmap, <-chan error) { +// +//} func (isdf *InterStepDataForward) ackRoutine(ctx context.Context, ackMsgChan <-chan *isb.ReadMessage, wg *sync.WaitGroup) { defer wg.Done() @@ -647,52 +582,6 @@ func (isdf *InterStepDataForward) writeToBuffer(ctx context.Context, toBufferPar return writeOffsets, nil } -//// concurrentApplyUDF applies the map UDF based on the request from the channel -//func (isdf *InterStepDataForward) concurrentApplyUDF(ctx context.Context, readMessagePair <-chan *readWriteMessagePair) { -// for message := range readMessagePair { -// start := time.Now() -// metrics.UDFReadMessagesCount.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() -// writeMessages, err := isdf.applyUDF(ctx, message.readMessage) -// 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(len(writeMessages))) -// // set the headers for the write messages -// for _, m := range writeMessages { -// m.Headers = message.readMessage.Headers -// } -// message.writeMessages = append(message.writeMessages, writeMessages...) -// message.udfError = err -// metrics.UDFProcessingTime.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())) -// } -//} - -//// applyUDF applies the map UDF and will block if there is any InternalErr. On the other hand, if this is a UserError -//// the skip flag is set. ShutDown flag will only if there is an InternalErr and ForceStop has been invoked. -//// The UserError retry will be done on the ApplyUDF. -//func (isdf *InterStepDataForward) applyUDF(ctx context.Context, readMessage *isb.ReadMessage) ([]*isb.WriteMessage, error) { -// for { -// writeMessages, err := isdf.flatmapUDF.ApplyMap(ctx, readMessage) -// if err != nil { -// isdf.opts.logger.Errorw("flatmapUDF.Apply error", zap.Error(err)) -// // TODO: implement retry with backoff etc. -// time.Sleep(isdf.opts.retryInterval) -// // keep retrying, I cannot think of a use case where a user could say, errors are fine :-) -// // as a platform we should not lose or corrupt data. -// // this does not mean we should prohibit this from a shutdown. -// if ok, _ := isdf.IsShuttingDown(); ok { -// isdf.opts.logger.Errorw("flatmapUDF.Apply, Stop called while stuck on an internal error", zap.Error(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 nil, err -// } -// continue -// } else { -// for index, m := range writeMessages { -// // add vertex name to the ID, since multiple vertices can publish to the same vertex and we need uniqueness across them -// m.ID = fmt.Sprintf("%s-%s-%d", readMessage.ReadOffset.String(), isdf.vertexName, index) -// } -// return writeMessages, 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 diff --git a/pkg/flatmap/rpc/grpc_flatmap.go b/pkg/flatmap/rpc/grpc_flatmap.go index 3814bba83b..7e6f96ff95 100644 --- a/pkg/flatmap/rpc/grpc_flatmap.go +++ b/pkg/flatmap/rpc/grpc_flatmap.go @@ -2,7 +2,6 @@ package rpc import ( "fmt" - "log" "time" flatmappb "github.com/numaproj/numaflow-go/pkg/apis/proto/flatmap/v1" @@ -59,63 +58,99 @@ func (u *GRPCBasedFlatmap) WaitUntilReady(ctx context.Context) error { } } +// 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 context.Context, messageStream []*isb.ReadMessage, responseCh chan<- *types.ResponseFlatmap) <-chan error { - var ( - errCh = make(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) + + // 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) - defer close(flatmapRequests) - // invoke the AsyncReduceFn method with requestCh channel and send the result to responseCh channel - // and any error to errCh channel + // 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): check this close here + // 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(responseCh) + // invoke the MapFn from the gRPC client for a stream of input requests + // resultCh -> chan to read responses streamed back + // reduceErrCh -> chan for reading any errors encountered during gRPC resultCh, reduceErrCh := 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 { - //// if the resultCh channel is closed, close the responseCh return } - resp, remove := u.parseMapResponse(result) + 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 with no data field (only EOR = true), we do not + // need to send it forward to the responseCh. if remove { - u.tracker.RemoveRequest(resp.Uid) + u.tracker.RemoveRequest(uid) continue } + // Forward the received response to the channel responseCh <- resp - log.Println("MYDEBUG: sending to writeCh", resp.Uid) case err := <-reduceErrCh: - // ctx.Done() event will be handled by the AsyncReduceFn method - // so we don't need a separate case for ctx.Done() here + // 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): check error here + // TODO(stream): Should we stop processing further in this case then + //return } } } }() - for _, req := range messageStream { - d := u.createMapRequest(req) - log.Print("MYDEBUG: TRYING TO SEND HERE", d.Uuid) - flatmapRequests <- d - } - + // 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) + flatmapRequests <- d + } + }() return errCh } +// createMapRequest takes a isb.ReadMessage and returns proto MapRequest func (u *GRPCBasedFlatmap) createMapRequest(msg *isb.ReadMessage) *flatmappb.MapRequest { 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, @@ -126,32 +161,33 @@ func (u *GRPCBasedFlatmap) createMapRequest(msg *isb.ReadMessage) *flatmappb.Map } return d } -func (u *GRPCBasedFlatmap) parseMapResponse(resp *flatmappb.MapResponse) (parsedResp *types.ResponseFlatmap, requestDone bool) { + +// 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() + uid = result.GetUuid() parentRequest, ok := u.tracker.GetRequest(uid) - // TODO(stream): check what should be path for !ok + // 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 { - // u.tracker.NewResponse(uid) - // idx = 1 } // Request has completed remove from the tracker module if eor == true { - return nil, true + return nil, true, uid } - //idx, present := u.tracker.GetIdx(uid) - //if !present { - - //} keys := result.GetKeys() taggedMessage := &isb.WriteMessage{ Message: isb.Message{ Header: isb.Header{ MessageInfo: parentRequest.MessageInfo, - // TODO(stream): Check what will be the unique ID to use here - //msgId := fmt.Sprintf("%s-%d-%s-%d", u.vertexName, u.vertexReplica, partitionID.String(), index) - //ID: fmt.Sprintf("%s-%s-%d", dataMessages[0].ReadOffset.String(), isdf.vertexName, msgIndex) + // TODO(stream): IMPORTANT Check what will be the unique ID to use here + // 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. ID: fmt.Sprintf("%s-%d", parentRequest.ReadOffset.String(), u.idx), Keys: keys, }, @@ -159,30 +195,25 @@ func (u *GRPCBasedFlatmap) parseMapResponse(resp *flatmappb.MapResponse) (parsed Payload: result.GetValue(), }, }, - Tags: result.GetTags(), } - u.idx += 1 - //u.tracker.IncrementRespIdx(uid) return &types.ResponseFlatmap{ ParentMessage: parentRequest, Uid: uid, RespMessage: taggedMessage, - }, false + }, false, uid } -// convertToUdfError converts the error returned by the reduceFn to ApplyUDFErr +// convertToUdfError converts the error returned by the MapFn to ApplyUDFErr func convertToUdfError(err error) error { - // if any error happens in reduce - // will exit and restart the numa container udfErr, _ := sdkerr.FromError(err) switch udfErr.ErrorKind() { case sdkerr.Retryable: - // TODO: currently we don't handle retryable errors for reduce + // TODO: currently we don't handle retryable errors yet return &ApplyUDFErr{ UserUDFErr: false, - Message: fmt.Sprintf("gRPC client.ReduceFn failed, %s", err), + Message: fmt.Sprintf("gRPC client.MapFn failed, %s", err), InternalErr: InternalErr{ Flag: true, MainCarDown: false, @@ -191,7 +222,7 @@ func convertToUdfError(err error) error { case sdkerr.NonRetryable: return &ApplyUDFErr{ UserUDFErr: false, - Message: fmt.Sprintf("gRPC client.ReduceFn failed, %s", err), + Message: fmt.Sprintf("gRPC client.MapFn failed, %s", err), InternalErr: InternalErr{ Flag: true, MainCarDown: false, @@ -209,7 +240,7 @@ func convertToUdfError(err error) error { default: return &ApplyUDFErr{ UserUDFErr: false, - Message: fmt.Sprintf("gRPC client.ReduceFn failed, %s", err), + 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 index 29a73d4c4b..962b6abb12 100644 --- a/pkg/flatmap/tracker/tracker.go +++ b/pkg/flatmap/tracker/tracker.go @@ -71,6 +71,7 @@ func (t *Tracker) RemoveRequest(id string) { func (t *Tracker) PrintAll() { t.lock.RLock() + defer t.lock.RUnlock() for k, v := range t.m { log.Println("MYDEBUG: MAP VALS", k, " ", v) } diff --git a/pkg/flatmap/types/types.go b/pkg/flatmap/types/types.go index e0a08a1228..19b5a0a51d 100644 --- a/pkg/flatmap/types/types.go +++ b/pkg/flatmap/types/types.go @@ -3,6 +3,8 @@ package types import "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 diff --git a/pkg/sdkclient/flatmapper/client.go b/pkg/sdkclient/flatmapper/client.go index d0c556a3b2..3d6e08b3b6 100644 --- a/pkg/sdkclient/flatmapper/client.go +++ b/pkg/sdkclient/flatmapper/client.go @@ -38,71 +38,91 @@ type client struct { 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 reciving the responses func (c client) MapFn(ctx context.Context, datumStreamCh <-chan *flatmappb.MapRequest) (<-chan *flatmappb.MapResponse, <-chan error) { var ( - errCh = make(chan error) + // 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) ) - log.Println("MYDEBUG: WHATS HAPPENING", len(datumStreamCh), datumStreamCh) - - // stream the messages to server + // MapFn is a bidirectional RPC + // We get a Flatmap_MapFnClient interface over which we can send the requests, + // receive the responses asynchronously. 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 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() - //if err == io.EOF { - // return - //} + // check if this is EOF error, which indicates that no more messages left to process on the + // stream, in such a case we return without any error if errors.Is(err, io.EOF) { + log.Println("MYDEBUG: ERROR GOT EOF", err) // skip selection on nil channel //errCh = nil //close(responseCh) return } + // If this is some other error, propagate it to error channel, + // also close the response channel to indicate no more messages being read errSDK := sdkerr.ToUDFErr("c.grpcClt.MapStreamFn", err) if errSDK != nil { log.Println("MYDEBUG: ERROR in recv", err, errSDK) errCh <- errSDK return } - log.Println("MYDEBUG: GOT IT FROM GRPC", resp.Result.Uuid) + //log.Println("MYDEBUG: GOT IT FROM GRPC", resp.Result.Uuid) responseCh <- resp } } }() - // Read from the inputStream and send messages - for inputMsg := range datumStreamCh { - log.Println("MYDEBUG: Sending to grpc", inputMsg.Uuid) - err = stream.Send(inputMsg) - if err != nil { + // 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 + } + //log.Println("MYDEBUG: SENT TO GRPC", inputMsg.GetUuid()) + } + sendErr := stream.CloseSend() + if sendErr != nil && !errors.Is(sendErr, io.EOF) { go func(sErr error) { - errCh <- sdkerr.ToUDFErr("c.grpcClt.MapFn", sErr) - }(err) - break + errCh <- sdkerr.ToUDFErr("c.grpcClt.MapFn stream.CloseSend()", sErr) + }(sendErr) } - //if err != nil { - // errCh <- sdkerr.ToUDFErr("MapFn stream.Send()", sendErr) - // return - //} - } - sendErr := stream.CloseSend() - if sendErr != nil && !errors.Is(sendErr, io.EOF) { - go func(sErr error) { - errCh <- sdkerr.ToUDFErr("c.grpcClt.MapFn stream.CloseSend()", sErr) - }(sendErr) - } - + }() return responseCh, errCh } From d81242de2b98794fd843a984394b410a2532da30 Mon Sep 17 00:00:00 2001 From: Sidhant Kohli Date: Mon, 3 Jun 2024 18:30:29 -0700 Subject: [PATCH 04/14] working commit Signed-off-by: Sidhant Kohli --- pkg/flatmap/forward/forward.go | 152 +++++++++++++++++------------ pkg/flatmap/rpc/grpc_flatmap.go | 13 ++- pkg/flatmap/tracker/tracker.go | 9 +- pkg/flatmap/types/types.go | 4 +- pkg/sdkclient/flatmapper/client.go | 2 +- 5 files changed, 106 insertions(+), 74 deletions(-) diff --git a/pkg/flatmap/forward/forward.go b/pkg/flatmap/forward/forward.go index 3df9629f3b..04ac489213 100644 --- a/pkg/flatmap/forward/forward.go +++ b/pkg/flatmap/forward/forward.go @@ -268,72 +268,90 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { // 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. - // writeChan is the channel on which the responses from the UDF are sent, and then this is consumed by - // the writer + // 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 done, to indicate that no further processing is required. - // TODO(stream): on an error, we should close the writeChan, so that the writer can stop processing. + // TODO(stream): on an error, we should close the udfRespCh, so that the writer can stop processing. // 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 - writeChan := make(chan *types.ResponseFlatmap) + udfRespCh := make(chan *types.ResponseFlatmap) // send the input messages for processing // The error channel is used to signal any errors that might have occurred during the UDF processing. - // TODO(stream): add the error handler for graceful handling - udfErrorCh := isdf.flatmapUDF.ApplyMap(ctx, dataMessages, writeChan) + // TODO(stream): add the error handler for graceful handling, check where should this be consumed + udfErrorCh := isdf.flatmapUDF.ApplyMap(ctx, dataMessages, udfRespCh) + // TODO(stream): got a error from the UDF, handle this gracefully. + go func() { + select { + 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 + } + if udfErr != nil { + isdf.opts.logger.Error("MYDEBUG: error during UDF processing", zap.Error(udfErr)) + // TODO(stream): got a error from the UDF + //isdf.opts.logger.Panic("Got an error while invoking ApplyMap", zap.Error(udfErr)) + } + } + }() // TODO(stream): Publish the watermark for these writers - ackChan := isdf.invokeWriter(ctx, writeChan, udfErrorCh) - done := isdf.invokeAck(ctx, ackChan) + // 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 + done := isdf.invokeAck(ctx, toAckChan) + // Wait until the acking has been completed <-done + //TODO(stream): we should ack the non data messages as well? - // 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 - //err = isdf.ackFromBuffer(ctx, readOffsets) - //// implicit return for posterity :-) - //if err != nil { - // isdf.opts.logger.Errorw("Failed to ack from buffer", zap.Error(err)) - // 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()}).Add(float64(len(readOffsets))) - // return - //} - //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(readOffsets))) + //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? - // 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("concurrent applyUDF completed", zap.Int("concurrency", isdf.opts.udfConcurrency), zap.Duration("took", time.Since(start))) + isdf.opts.logger.Debugw("forwardAChunk completed", zap.Int("concurrency", isdf.opts.udfConcurrency), zap.Duration("took", time.Since(start))) } -// -//// invokeUDF reads requests from the supplied PBQ, invokes the UDF to gets the response and writes the response to the -//// main channel. -//func (isdf *InterStepDataForward) invokeUDF(ctx context.Context, dataMessages []*isb.ReadMessage) (<-chan *types.ResponseFlatmap, <-chan error) { -// -//} - +// 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 *isb.ReadMessage, wg *sync.WaitGroup) { defer wg.Done() -forwardLoop: +ackLoop: for { select { case <-ctx.Done(): - break forwardLoop - + break ackLoop case response, ok := <-ackMsgChan: if !ok { - break forwardLoop + break ackLoop } ackMessages := []isb.Offset{response.ReadOffset} if err := isdf.ackFromBuffer(ctx, ackMessages); err != nil { isdf.opts.logger.Error("MYDEBUG: ERROR IN ACK ", zap.Error(err)) - // TODO(stream): we have retried in the ackFromBuffer, should we trigger - // shutdown here then? + // TODO(stream): we have retried in the ackFromBuffer, should we trigger drain here then? + 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.opts.logger.Info("MYDEBUG: NEW GOT TO ACK THIS ", response.ReadOffset) + 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()}).Inc() } } } + +// 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 *isb.ReadMessage) (doneChan chan struct{}) { logger := isdf.opts.logger logger.Info("MYDEBUG: NO WG ACK ROUTINE ", isdf.opts.readBatchSize) @@ -350,70 +368,78 @@ func (isdf *InterStepDataForward) invokeAck(ctx context.Context, ackMsgChan <-ch return doneChan } -func (isdf *InterStepDataForward) writeRoutine(ctx context.Context, writeMessageCh <-chan *types.ResponseFlatmap, errCh <-chan error, ackChan chan<- *isb.ReadMessage, wg *sync.WaitGroup) { +// 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 *types.ResponseFlatmap, ackChan chan<- *isb.ReadMessage, wg *sync.WaitGroup) { defer wg.Done() outerLoop: for { select { - case err := <-errCh: - if errors.Is(err, context.Canceled) || ctx.Err() != nil { - isdf.opts.logger.Infow("Context is canceled, stopping the writer", zap.Error(err)) - return - } - if err != nil { - isdf.opts.logger.Error("MYDEBUG: NEW ERROR IN ERROR CH WRITE", zap.Error(err)) - // TODO(stream): got a error from the UDF, time to start the shutdown then or retry - //isdf.opts.logger.Panic("Got an error while invoking ApplyMap", zap.Error(err)) - } - case response, ok := <-writeMessageCh: + case response, ok := <-udfRespCh: + // if the channel is closed, exit from the routine if !ok { break outerLoop } - //isdf.opts.logger.Info("MYDEBUG: NEW SENDING TO ACK ", response.ParentMessage.ReadOffset, " ", response.Uid) + // 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 response.AckIt { + ackChan <- response.ParentMessage + 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 := response.RespMessage - writeMessages := []*isb.WriteMessage{writeMessage} - if err := isdf.forwardToBuffers(ctx, &writeMessages, response.ParentMessage, messageToStep); err != nil { - // // TODO(stream): mark as no ack directly or retry + if err := isdf.forwardToBuffers(ctx, writeMessage, response.ParentMessage, messageToStep); err != nil { + // TODO(stream): mark as no ack directly or retry isdf.opts.logger.Error("MYDEBUG: NEW ERROR IN WRITE", zap.Error(err)) } - 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(1)) - ackChan <- response.ParentMessage } } } -func (isdf *InterStepDataForward) invokeWriter(ctx context.Context, writeMessageCh <-chan *types.ResponseFlatmap, errCh <-chan error) <-chan *isb.ReadMessage { +// 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 *types.ResponseFlatmap) <-chan *isb.ReadMessage { + // 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 *isb.ReadMessage) 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, errCh, ackChan, &group) + go isdf.writeRoutine(ctx, writeMessageCh, ackChan, &group) } + // wait untill 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 len(*writeMessages) == 0 { +func (isdf *InterStepDataForward) forwardToBuffers(ctx context.Context, writeMessages *isb.WriteMessage, readMessage *isb.ReadMessage, messageToStep map[string][]isb.Message) error { + if writeMessages == nil { return nil } - for _, message := range *writeMessages { - if err := isdf.whereToStep(message, messageToStep, readMessage); err != nil { - isdf.opts.logger.Errorw("failed in whereToStep", zap.Error(err)) - return err - } + 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 { diff --git a/pkg/flatmap/rpc/grpc_flatmap.go b/pkg/flatmap/rpc/grpc_flatmap.go index 7e6f96ff95..08a2c00ff8 100644 --- a/pkg/flatmap/rpc/grpc_flatmap.go +++ b/pkg/flatmap/rpc/grpc_flatmap.go @@ -100,11 +100,10 @@ func (u *GRPCBasedFlatmap) ApplyMap(ctx context.Context, messageStream []*isb.Re } 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 with no data field (only EOR = true), we do not - // need to send it forward to the responseCh. + // 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) - continue } // Forward the received response to the channel responseCh <- resp @@ -177,7 +176,12 @@ func (u *GRPCBasedFlatmap) parseMapResponse(resp *flatmappb.MapResponse) (parsed } // Request has completed remove from the tracker module if eor == true { - return nil, true, uid + return &types.ResponseFlatmap{ + ParentMessage: parentRequest, + Uid: uid, + RespMessage: nil, + AckIt: true, + }, true, uid } keys := result.GetKeys() taggedMessage := &isb.WriteMessage{ @@ -202,6 +206,7 @@ func (u *GRPCBasedFlatmap) parseMapResponse(resp *flatmappb.MapResponse) (parsed ParentMessage: parentRequest, Uid: uid, RespMessage: taggedMessage, + AckIt: false, }, false, uid } diff --git a/pkg/flatmap/tracker/tracker.go b/pkg/flatmap/tracker/tracker.go index 962b6abb12..ce543943e4 100644 --- a/pkg/flatmap/tracker/tracker.go +++ b/pkg/flatmap/tracker/tracker.go @@ -11,7 +11,7 @@ import ( type Tracker struct { requestMap sync.Map - // TODO(stream): check if this will be inefficient, and there is a better way for this + // TODO(stream): check perf between sync.map and mutex+map responseIdx sync.Map lock sync.RWMutex m map[string]*isb.ReadMessage @@ -19,9 +19,9 @@ type Tracker struct { func NewTracker() *Tracker { return &Tracker{ - requestMap: sync.Map{}, - responseIdx: sync.Map{}, - m: make(map[string]*isb.ReadMessage), + //requestMap: sync.Map{}, + //responseIdx: sync.Map{}, + m: make(map[string]*isb.ReadMessage), } } @@ -31,6 +31,7 @@ func GetNewId() 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) diff --git a/pkg/flatmap/types/types.go b/pkg/flatmap/types/types.go index 19b5a0a51d..d6f42310c3 100644 --- a/pkg/flatmap/types/types.go +++ b/pkg/flatmap/types/types.go @@ -3,12 +3,12 @@ package types import "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 +// We might just need only a few details out of these type ResponseFlatmap struct { ParentMessage *isb.ReadMessage Uid string RespMessage *isb.WriteMessage + AckIt bool } type WriteMsgFlatmap struct { diff --git a/pkg/sdkclient/flatmapper/client.go b/pkg/sdkclient/flatmapper/client.go index 3d6e08b3b6..0a4f0b8997 100644 --- a/pkg/sdkclient/flatmapper/client.go +++ b/pkg/sdkclient/flatmapper/client.go @@ -86,7 +86,7 @@ func (c client) MapFn(ctx context.Context, datumStreamCh <-chan *flatmappb.MapRe if errors.Is(err, io.EOF) { log.Println("MYDEBUG: ERROR GOT EOF", err) // skip selection on nil channel - //errCh = nil + errCh = nil //close(responseCh) return } From 157b37569326fba8bb1d78b87c0da352ebc57abf Mon Sep 17 00:00:00 2001 From: Sidhant Kohli Date: Tue, 4 Jun 2024 17:30:08 -0700 Subject: [PATCH 05/14] eof fix Signed-off-by: Sidhant Kohli --- pkg/flatmap/forward/forward.go | 20 ++++++++++++++------ pkg/flatmap/rpc/grpc_flatmap.go | 6 ++++++ pkg/sdkclient/flatmapper/client.go | 27 +++++++++++++-------------- 3 files changed, 33 insertions(+), 20 deletions(-) diff --git a/pkg/flatmap/forward/forward.go b/pkg/flatmap/forward/forward.go index 04ac489213..ac8ac5fb17 100644 --- a/pkg/flatmap/forward/forward.go +++ b/pkg/flatmap/forward/forward.go @@ -277,23 +277,31 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { // either as ack/no ack, also to check the responsibility for close udfRespCh := make(chan *types.ResponseFlatmap) - // send the input messages for processing - // The error channel is used to signal any errors that might have occurred during the UDF processing. - // TODO(stream): add the error handler for graceful handling, check where should this be consumed + // Send the input messages for processing + // The error channel returned is used to signal any errors that might have occurred during the UDF processing. udfErrorCh := isdf.flatmapUDF.ApplyMap(ctx, dataMessages, udfRespCh) // TODO(stream): got a error from the UDF, handle this gracefully. go func() { select { case udfErr := <-udfErrorCh: - // TODO(stream): when context is cancelled, do we want to do some different handling + // 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 + // 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.Error("MYDEBUG: error during UDF processing", zap.Error(udfErr)) // TODO(stream): got a error from the UDF - //isdf.opts.logger.Panic("Got an error while invoking ApplyMap", zap.Error(udfErr)) + isdf.opts.logger.Panic("Got an error while invoking ApplyMap", zap.Error(udfErr)) } } }() @@ -310,7 +318,7 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { // We keep reading on the toAckChan to consume the messages and then ack them to the prev buffer done := isdf.invokeAck(ctx, toAckChan) - // Wait until the acking has been completed + // Wait until the Acking has been completed <-done //TODO(stream): we should ack the non data messages as well? diff --git a/pkg/flatmap/rpc/grpc_flatmap.go b/pkg/flatmap/rpc/grpc_flatmap.go index 08a2c00ff8..f128dcad37 100644 --- a/pkg/flatmap/rpc/grpc_flatmap.go +++ b/pkg/flatmap/rpc/grpc_flatmap.go @@ -95,6 +95,7 @@ func (u *GRPCBasedFlatmap) ApplyMap(ctx context.Context, messageStream []*isb.Re 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 + // gRPC has no more messages to send. Hence, we can return from here if !ok || result == nil { return } @@ -253,3 +254,8 @@ func convertToUdfError(err error) error { } } } + +func (u *GRPCBasedFlatmap) RemoveFromTracker(uid string) { + u.tracker.RemoveRequest(uid) + +} diff --git a/pkg/sdkclient/flatmapper/client.go b/pkg/sdkclient/flatmapper/client.go index 0a4f0b8997..b79a2b97e5 100644 --- a/pkg/sdkclient/flatmapper/client.go +++ b/pkg/sdkclient/flatmapper/client.go @@ -20,7 +20,6 @@ import ( "context" "errors" "io" - "log" flatmappb "github.com/numaproj/numaflow-go/pkg/apis/proto/flatmap/v1" "github.com/numaproj/numaflow-go/pkg/info" @@ -42,7 +41,7 @@ type client struct { // 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 reciving the responses +// 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 @@ -71,6 +70,9 @@ func (c client) MapFn(ctx context.Context, datumStreamCh <-chan *flatmappb.MapRe // 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 { @@ -80,21 +82,16 @@ func (c client) MapFn(ctx context.Context, datumStreamCh <-chan *flatmappb.MapRe return default: var resp *flatmappb.MapResponse - resp, err = stream.Recv() - // check if this is EOF error, which indicates that no more messages left to process on the - // stream, in such a case we return without any error + 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) { - log.Println("MYDEBUG: ERROR GOT EOF", err) - // skip selection on nil channel - errCh = nil - //close(responseCh) return } // If this is some other error, propagate it to error channel, // also close the response channel to indicate no more messages being read - errSDK := sdkerr.ToUDFErr("c.grpcClt.MapStreamFn", err) + errSDK := sdkerr.ToUDFErr("flatmap c.grpcClt.MapFn", err) if errSDK != nil { - log.Println("MYDEBUG: ERROR in recv", err, errSDK) errCh <- errSDK return } @@ -107,19 +104,21 @@ func (c client) MapFn(ctx context.Context, datumStreamCh <-chan *flatmappb.MapRe // 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) + err := stream.Send(inputMsg) if err != nil { go func(sErr error) { errCh <- sdkerr.ToUDFErr("c.grpcClt.MapFn", sErr) }(err) break } - //log.Println("MYDEBUG: SENT TO GRPC", inputMsg.GetUuid()) } + // 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("c.grpcClt.MapFn stream.CloseSend()", sErr) + errCh <- sdkerr.ToUDFErr("flatmap c.grpcClt.MapFn stream.CloseSend()", sErr) }(sendErr) } }() From 2a7a8b4d38be6328256cabd53abf294d59a845fa Mon Sep 17 00:00:00 2001 From: Sidhant Kohli Date: Tue, 4 Jun 2024 17:32:13 -0700 Subject: [PATCH 06/14] comments Signed-off-by: Sidhant Kohli --- pkg/flatmap/rpc/grpc_flatmap.go | 5 ----- pkg/sdkclient/flatmapper/client.go | 3 +-- 2 files changed, 1 insertion(+), 7 deletions(-) diff --git a/pkg/flatmap/rpc/grpc_flatmap.go b/pkg/flatmap/rpc/grpc_flatmap.go index f128dcad37..149e4d6fa6 100644 --- a/pkg/flatmap/rpc/grpc_flatmap.go +++ b/pkg/flatmap/rpc/grpc_flatmap.go @@ -254,8 +254,3 @@ func convertToUdfError(err error) error { } } } - -func (u *GRPCBasedFlatmap) RemoveFromTracker(uid string) { - u.tracker.RemoveRequest(uid) - -} diff --git a/pkg/sdkclient/flatmapper/client.go b/pkg/sdkclient/flatmapper/client.go index b79a2b97e5..b21c1ccf02 100644 --- a/pkg/sdkclient/flatmapper/client.go +++ b/pkg/sdkclient/flatmapper/client.go @@ -89,13 +89,12 @@ func (c client) MapFn(ctx context.Context, datumStreamCh <-chan *flatmappb.MapRe return } // If this is some other error, propagate it to error channel, - // also close the response channel to indicate no more messages being read + // 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.Println("MYDEBUG: GOT IT FROM GRPC", resp.Result.Uuid) responseCh <- resp } } From 61ba68a1b6061cf29e5a0d32f8e0a25e028768bb Mon Sep 17 00:00:00 2001 From: Sidhant Kohli Date: Wed, 5 Jun 2024 11:40:55 -0700 Subject: [PATCH 07/14] add unique msg id Signed-off-by: Sidhant Kohli --- pkg/flatmap/flatmap_udf.go | 2 +- pkg/flatmap/forward/forward.go | 7 ------- pkg/flatmap/rpc/grpc_flatmap.go | 33 +++++++++++++++++++++------------ 3 files changed, 22 insertions(+), 20 deletions(-) diff --git a/pkg/flatmap/flatmap_udf.go b/pkg/flatmap/flatmap_udf.go index 4128ab98a5..576a2556b4 100644 --- a/pkg/flatmap/flatmap_udf.go +++ b/pkg/flatmap/flatmap_udf.go @@ -134,7 +134,7 @@ func (u *FlatmapUDFProcessor) Start(ctx context.Context) error { if err != nil { return fmt.Errorf("failed to create map client, %w", err) } - flatmapHandler = rpc.NewUDSgRPCBasedFlatmap(flatmapClient, int(*u.VertexInstance.Vertex.Spec.Limits.ReadBatchSize)) + flatmapHandler = rpc.NewUDSgRPCBasedFlatmap(flatmapClient, u.VertexInstance.Vertex.Spec.Name) // Readiness check if err := flatmapHandler.WaitUntilReady(ctx); err != nil { diff --git a/pkg/flatmap/forward/forward.go b/pkg/flatmap/forward/forward.go index ac8ac5fb17..4265e7caf1 100644 --- a/pkg/flatmap/forward/forward.go +++ b/pkg/flatmap/forward/forward.go @@ -167,13 +167,6 @@ func (isdf *InterStepDataForward) Start() <-chan struct{} { return stopped } -// readWriteMessagePair represents a read message and its processed (via map UDF) write messages. -type readWriteMessagePair struct { - readMessage *isb.ReadMessage - writeMessages []*isb.WriteMessage - udfError error -} - // 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. diff --git a/pkg/flatmap/rpc/grpc_flatmap.go b/pkg/flatmap/rpc/grpc_flatmap.go index 149e4d6fa6..b9fa0f0bb3 100644 --- a/pkg/flatmap/rpc/grpc_flatmap.go +++ b/pkg/flatmap/rpc/grpc_flatmap.go @@ -20,14 +20,13 @@ import ( // 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 - readBatchSize int - idx int + client flatmapper.Client + tracker *tracker.Tracker + vertexName string } -func NewUDSgRPCBasedFlatmap(client flatmapper.Client, batchSize int) *GRPCBasedFlatmap { - return &GRPCBasedFlatmap{client: client, tracker: tracker.NewTracker(), readBatchSize: batchSize} +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. @@ -189,11 +188,22 @@ func (u *GRPCBasedFlatmap) parseMapResponse(resp *flatmappb.MapResponse) (parsed Message: isb.Message{ Header: isb.Header{ MessageInfo: parentRequest.MessageInfo, - // TODO(stream): IMPORTANT Check what will be the unique ID to use here - // 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. - ID: fmt.Sprintf("%s-%d", parentRequest.ReadOffset.String(), u.idx), + // 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{ @@ -202,7 +212,6 @@ func (u *GRPCBasedFlatmap) parseMapResponse(resp *flatmappb.MapResponse) (parsed }, Tags: result.GetTags(), } - u.idx += 1 return &types.ResponseFlatmap{ ParentMessage: parentRequest, Uid: uid, From 6f02ee0470f5cebb7fe4a85141234703b2bfc6da Mon Sep 17 00:00:00 2001 From: Sidhant Kohli Date: Wed, 5 Jun 2024 11:50:14 -0700 Subject: [PATCH 08/14] add unique msg id Signed-off-by: Sidhant Kohli --- pkg/flatmap/forward/forward.go | 7 +++---- pkg/flatmap/rpc/grpc_flatmap.go | 2 +- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/pkg/flatmap/forward/forward.go b/pkg/flatmap/forward/forward.go index 4265e7caf1..d4057645d2 100644 --- a/pkg/flatmap/forward/forward.go +++ b/pkg/flatmap/forward/forward.go @@ -273,7 +273,7 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { // Send the input messages for processing // The error channel returned is used to signal any errors that might have occurred during the UDF processing. udfErrorCh := isdf.flatmapUDF.ApplyMap(ctx, dataMessages, udfRespCh) - // TODO(stream): got a error from the UDF, handle this gracefully. + // TODO(stream): got an error from the UDF, handle this gracefully. go func() { select { case udfErr := <-udfErrorCh: @@ -287,13 +287,12 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { // 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 + // 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.Error("MYDEBUG: error during UDF processing", zap.Error(udfErr)) - // TODO(stream): got a error from the UDF isdf.opts.logger.Panic("Got an error while invoking ApplyMap", zap.Error(udfErr)) } } diff --git a/pkg/flatmap/rpc/grpc_flatmap.go b/pkg/flatmap/rpc/grpc_flatmap.go index b9fa0f0bb3..e686aef8e5 100644 --- a/pkg/flatmap/rpc/grpc_flatmap.go +++ b/pkg/flatmap/rpc/grpc_flatmap.go @@ -202,7 +202,7 @@ func (u *GRPCBasedFlatmap) parseMapResponse(resp *flatmappb.MapResponse) (parsed // // 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 + // 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, }, From 9f1d78e426c1201e6f731b8c3a113aed7e42d5f0 Mon Sep 17 00:00:00 2001 From: Sidhant Kohli Date: Wed, 5 Jun 2024 15:15:49 -0700 Subject: [PATCH 09/14] comments Signed-off-by: Sidhant Kohli --- pkg/flatmap/forward/forward.go | 36 +++++++++++++++++++----------- pkg/flatmap/rpc/grpc_flatmap.go | 2 +- pkg/sdkclient/flatmapper/client.go | 3 +++ 3 files changed, 27 insertions(+), 14 deletions(-) diff --git a/pkg/flatmap/forward/forward.go b/pkg/flatmap/forward/forward.go index d4057645d2..a87bd7def8 100644 --- a/pkg/flatmap/forward/forward.go +++ b/pkg/flatmap/forward/forward.go @@ -194,9 +194,9 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { } // 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 done - //// 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. + // 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 @@ -225,7 +225,6 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { // } // return //} - // TODO(stream): remove once idle watermark is done if len(readMessages) == 0 { return } @@ -234,12 +233,15 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { // We send only the dataMessages to the UDF for processing, for the non data messages, // we just need to ack? var dataMessages = make([]*isb.ReadMessage, 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 { dataMessages = append(dataMessages, m) + } else { + ctrlMessageOffsets = append(ctrlMessageOffsets, m.ReadOffset) } } @@ -264,8 +266,7 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { // 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 done, to indicate that no further processing is required. - // TODO(stream): on an error, we should close the udfRespCh, so that the writer can stop processing. + // 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 *types.ResponseFlatmap) @@ -308,14 +309,22 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { // 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 - done := isdf.invokeAck(ctx, toAckChan) + ackDone := isdf.invokeAck(ctx, toAckChan) // Wait until the Acking has been completed - <-done - //TODO(stream): we should ack the non data messages as well? + <-ackDone + + //TODO(stream): is it fine if we ack the ctrlMessageOffsets in the end? + // ack the control messages, also based on some error do + 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? + // or check if any messages are left in that which have not been processed and those can be noAcked? isdf.opts.logger.Debugw("forwardAChunk completed", zap.Int("concurrency", isdf.opts.udfConcurrency), zap.Duration("took", time.Since(start))) } @@ -384,7 +393,7 @@ outerLoop: break outerLoop } // 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 + // a given message. In this case send the parent request for Acking and continue if response.AckIt { ackChan <- response.ParentMessage continue @@ -397,7 +406,8 @@ outerLoop: } writeMessage := response.RespMessage if err := isdf.forwardToBuffers(ctx, writeMessage, response.ParentMessage, messageToStep); err != nil { - // TODO(stream): mark as no ack directly or retry + // As we have re-tried already to forward to the buffer, we should not be trying it again. + // But what if we isdf.opts.logger.Error("MYDEBUG: NEW ERROR IN WRITE", zap.Error(err)) } 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(1)) @@ -426,7 +436,7 @@ func (isdf *InterStepDataForward) invokeWriter(ctx context.Context, writeMessage group.Add(1) go isdf.writeRoutine(ctx, writeMessageCh, ackChan, &group) } - // wait untill all the writeRoutines are done + // wait until all the writeRoutines are done group.Wait() }() return ackChan diff --git a/pkg/flatmap/rpc/grpc_flatmap.go b/pkg/flatmap/rpc/grpc_flatmap.go index e686aef8e5..e82870602b 100644 --- a/pkg/flatmap/rpc/grpc_flatmap.go +++ b/pkg/flatmap/rpc/grpc_flatmap.go @@ -121,7 +121,7 @@ func (u *GRPCBasedFlatmap) ApplyMap(ctx context.Context, messageStream []*isb.Re // 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 + return } } } diff --git a/pkg/sdkclient/flatmapper/client.go b/pkg/sdkclient/flatmapper/client.go index b21c1ccf02..3cc09ba4a1 100644 --- a/pkg/sdkclient/flatmapper/client.go +++ b/pkg/sdkclient/flatmapper/client.go @@ -56,6 +56,9 @@ func (c client) MapFn(ctx context.Context, datumStreamCh <-chan *flatmappb.MapRe // 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 From 1cb22d4af1b6af01139f845a51fa0e85dc334b06 Mon Sep 17 00:00:00 2001 From: Sidhant Kohli Date: Wed, 5 Jun 2024 18:19:19 -0700 Subject: [PATCH 10/14] add write/ack tracker Signed-off-by: Sidhant Kohli --- pkg/flatmap/forward/forward.go | 23 +++++++++++++++++++++-- pkg/flatmap/rpc/grpc_flatmap.go | 2 ++ pkg/flatmap/types/types.go | 1 + 3 files changed, 24 insertions(+), 2 deletions(-) diff --git a/pkg/flatmap/forward/forward.go b/pkg/flatmap/forward/forward.go index a87bd7def8..5c4aed6338 100644 --- a/pkg/flatmap/forward/forward.go +++ b/pkg/flatmap/forward/forward.go @@ -26,6 +26,7 @@ import ( "math" "strconv" "sync" + "sync/atomic" "time" "go.uber.org/zap" @@ -66,6 +67,7 @@ type InterStepDataForward struct { idleManager wmb.IdleManager // wmbChecker checks if the idle watermark is valid when the len(readMessage) is 0. wmbChecker wmb.WMBChecker + counters *sync.Map Shutdown } @@ -99,7 +101,8 @@ func NewInterStepDataForward(vertexInstance *dfv1.VertexInstance, fromStep isb.B Shutdown: Shutdown{ rwlock: new(sync.RWMutex), }, - opts: *options, + opts: *options, + counters: new(sync.Map), } // Add logger from parent ctx to child context. @@ -323,6 +326,8 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { } } + isdf.counters = new(sync.Map) + //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? @@ -394,7 +399,7 @@ outerLoop: } // 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 response.AckIt { + if response.AckIt && response.Total == 0 { ackChan <- response.ParentMessage continue } @@ -410,12 +415,26 @@ outerLoop: // But what if we isdf.opts.logger.Error("MYDEBUG: NEW ERROR IN WRITE", zap.Error(err)) } + // Updat the response counter for the given UUID (request) + idxNum := isdf.updateCounter(response.Uid) + // If the counter has reached the total number of responses expected, we can safely + // send the parent request for Acking + if idxNum == response.Total { + ackChan <- response.ParentMessage + } 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(1)) } } } +func (isdf *InterStepDataForward) updateCounter(key string) int64 { + val, _ := isdf.counters.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. diff --git a/pkg/flatmap/rpc/grpc_flatmap.go b/pkg/flatmap/rpc/grpc_flatmap.go index e82870602b..32f50e2884 100644 --- a/pkg/flatmap/rpc/grpc_flatmap.go +++ b/pkg/flatmap/rpc/grpc_flatmap.go @@ -181,6 +181,7 @@ func (u *GRPCBasedFlatmap) parseMapResponse(resp *flatmappb.MapResponse) (parsed Uid: uid, RespMessage: nil, AckIt: true, + Total: int64(resp.Result.Total), }, true, uid } keys := result.GetKeys() @@ -217,6 +218,7 @@ func (u *GRPCBasedFlatmap) parseMapResponse(resp *flatmappb.MapResponse) (parsed Uid: uid, RespMessage: taggedMessage, AckIt: false, + Total: int64(resp.Result.Total), }, false, uid } diff --git a/pkg/flatmap/types/types.go b/pkg/flatmap/types/types.go index d6f42310c3..1eb744b043 100644 --- a/pkg/flatmap/types/types.go +++ b/pkg/flatmap/types/types.go @@ -9,6 +9,7 @@ type ResponseFlatmap struct { Uid string RespMessage *isb.WriteMessage AckIt bool + Total int64 } type WriteMsgFlatmap struct { From a3e8e0374247c07a871027bbf1e2eab577feb0c3 Mon Sep 17 00:00:00 2001 From: Sidhant Kohli Date: Sun, 9 Jun 2024 22:12:08 -0700 Subject: [PATCH 11/14] update trackers Signed-off-by: Sidhant Kohli --- pkg/flatmap/forward/applier/mapper.go | 9 +- pkg/flatmap/forward/forward.go | 140 +++++++++++++++++++++----- pkg/flatmap/rpc/grpc_flatmap.go | 123 +++++++++++----------- pkg/flatmap/tracker/tracker.go | 39 +++++-- pkg/flatmap/types/types.go | 13 ++- pkg/sdkclient/flatmapper/client.go | 2 + 6 files changed, 228 insertions(+), 98 deletions(-) diff --git a/pkg/flatmap/forward/applier/mapper.go b/pkg/flatmap/forward/applier/mapper.go index cfb6ee2f10..b3cbbfe355 100644 --- a/pkg/flatmap/forward/applier/mapper.go +++ b/pkg/flatmap/forward/applier/mapper.go @@ -19,8 +19,9 @@ package applier import ( "context" + flatmappb "github.com/numaproj/numaflow-go/pkg/apis/proto/flatmap/v1" + "github.com/numaproj/numaflow/pkg/flatmap/types" - "github.com/numaproj/numaflow/pkg/isb" ) // FlatmapApplier applies the GRPCBasedMapUDF on the stream of read messages and gives back a new message. @@ -28,12 +29,12 @@ 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 []*isb.ReadMessage, writeChan chan<- *types.ResponseFlatmap) <-chan error + 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 []*isb.ReadMessage, writeChan chan<- *types.ResponseFlatmap) <-chan error +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 []*isb.ReadMessage, writeChan chan<- *types.ResponseFlatmap) <-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/forward.go b/pkg/flatmap/forward/forward.go index 5c4aed6338..e03ebd8e03 100644 --- a/pkg/flatmap/forward/forward.go +++ b/pkg/flatmap/forward/forward.go @@ -29,11 +29,13 @@ import ( "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" @@ -66,8 +68,9 @@ type InterStepDataForward struct { // 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 - counters *sync.Map + wmbChecker wmb.WMBChecker + responseTracker *sync.Map + requestTracker *tracker.Tracker Shutdown } @@ -101,8 +104,9 @@ func NewInterStepDataForward(vertexInstance *dfv1.VertexInstance, fromStep isb.B Shutdown: Shutdown{ rwlock: new(sync.RWMutex), }, - opts: *options, - counters: new(sync.Map), + opts: *options, + responseTracker: new(sync.Map), + requestTracker: tracker.NewTracker(), } // Add logger from parent ctx to child context. @@ -235,14 +239,17 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { // 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([]*isb.ReadMessage, 0, len(readMessages)) + 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 { - dataMessages = append(dataMessages, m) + newRequest := isdf.createNewRequest(m) + //isdf.opts.logger.Info("MYDEBUG NEW REQUEST ", newRequest.Uid) + dataMessages = append(dataMessages, newRequest) + } else { ctrlMessageOffsets = append(ctrlMessageOffsets, m.ReadOffset) } @@ -272,14 +279,19 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { // 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 *types.ResponseFlatmap) + 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. - udfErrorCh := isdf.flatmapUDF.ApplyMap(ctx, dataMessages, udfRespCh) + 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 { @@ -326,7 +338,14 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { } } - isdf.counters = new(sync.Map) + 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) + } + + isdf.responseTracker = new(sync.Map) //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? @@ -337,7 +356,7 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { // 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 *isb.ReadMessage, wg *sync.WaitGroup) { +func (isdf *InterStepDataForward) ackRoutine(ctx context.Context, ackMsgChan <-chan *types.RequestFlatmap, wg *sync.WaitGroup) { defer wg.Done() ackLoop: for { @@ -348,13 +367,17 @@ ackLoop: if !ok { break ackLoop } - ackMessages := []isb.Offset{response.ReadOffset} + //isdf.opts.logger.Info("MYDEBUG: GOT TO ACK ", response.Uid) + ackMessages := []isb.Offset{response.Request.ReadOffset} if err := isdf.ackFromBuffer(ctx, ackMessages); 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) 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()}).Inc() } } @@ -366,7 +389,7 @@ ackLoop: // 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 *isb.ReadMessage) (doneChan chan struct{}) { +func (isdf *InterStepDataForward) invokeAck(ctx context.Context, ackMsgChan <-chan *types.RequestFlatmap) (doneChan chan struct{}) { logger := isdf.opts.logger logger.Info("MYDEBUG: NO WG ACK ROUTINE ", isdf.opts.readBatchSize) doneChan = make(chan struct{}) @@ -387,7 +410,7 @@ func (isdf *InterStepDataForward) invokeAck(ctx context.Context, ackMsgChan <-ch // 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 *types.ResponseFlatmap, ackChan chan<- *isb.ReadMessage, wg *sync.WaitGroup) { +func (isdf *InterStepDataForward) writeRoutine(ctx context.Context, udfRespCh <-chan *flatmappb.MapResponse, ackChan chan<- *types.RequestFlatmap, wg *sync.WaitGroup) { defer wg.Done() outerLoop: for { @@ -397,10 +420,21 @@ outerLoop: if !ok { break outerLoop } + 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 response.AckIt && response.Total == 0 { - ackChan <- response.ParentMessage + // 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 @@ -409,18 +443,19 @@ outerLoop: // over allocating to have a predictable pattern messageToStep[toVertex] = make([]isb.Message, len(isdf.toBuffers[toVertex])) } - writeMessage := response.RespMessage - if err := isdf.forwardToBuffers(ctx, writeMessage, response.ParentMessage, messageToStep); err != nil { + 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. // But what if we isdf.opts.logger.Error("MYDEBUG: NEW ERROR IN WRITE", zap.Error(err)) } - // Updat the response counter for the given UUID (request) - idxNum := isdf.updateCounter(response.Uid) + // 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 == response.Total { - ackChan <- response.ParentMessage + if idxNum == int64(total) { + //isdf.opts.logger.Info("MYDEBUG: SENDING TO ACK ", response.Result.GetTotal()) + ackChan <- trackedRequest } 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(1)) } @@ -429,7 +464,7 @@ outerLoop: } func (isdf *InterStepDataForward) updateCounter(key string) int64 { - val, _ := isdf.counters.LoadOrStore(key, new(int64)) + val, _ := isdf.responseTracker.LoadOrStore(key, new(int64)) ptr := val.(*int64) newVal := atomic.AddInt64(ptr, 1) return newVal @@ -442,10 +477,10 @@ func (isdf *InterStepDataForward) updateCounter(key string) int64 { // 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 *types.ResponseFlatmap) <-chan *isb.ReadMessage { +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 *isb.ReadMessage) + ackChan := make(chan *types.RequestFlatmap) go func() { // close to indicate that no further messages left to ack defer close(ackChan) @@ -671,3 +706,60 @@ func errorArrayToMap(errs []error) map[string]int64 { } 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: fmt.Sprintf("%s-%s-%s", parentRequest.ReadOffset.String(), isdf.vertexName, result.GetIndex()), + Keys: keys, + }, + Body: isb.Body{ + Payload: result.GetValue(), + }, + }, + Tags: result.GetTags(), + } + return taggedMessage, false, total +} diff --git a/pkg/flatmap/rpc/grpc_flatmap.go b/pkg/flatmap/rpc/grpc_flatmap.go index 32f50e2884..1783072c34 100644 --- a/pkg/flatmap/rpc/grpc_flatmap.go +++ b/pkg/flatmap/rpc/grpc_flatmap.go @@ -1,6 +1,7 @@ package rpc import ( + context2 "context" "fmt" "time" @@ -11,7 +12,6 @@ import ( "github.com/numaproj/numaflow/pkg/flatmap/tracker" "github.com/numaproj/numaflow/pkg/flatmap/types" - "github.com/numaproj/numaflow/pkg/isb" sdkerr "github.com/numaproj/numaflow/pkg/sdkclient/error" "github.com/numaproj/numaflow/pkg/sdkclient/flatmapper" "github.com/numaproj/numaflow/pkg/shared/logging" @@ -60,10 +60,11 @@ func (u *GRPCBasedFlatmap) WaitUntilReady(ctx context.Context) error { // 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 context.Context, messageStream []*isb.ReadMessage, responseCh chan<- *types.ResponseFlatmap) <-chan error { +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? @@ -81,7 +82,7 @@ func (u *GRPCBasedFlatmap) ApplyMap(ctx context.Context, messageStream []*isb.Re // 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(responseCh) + defer close(doneChan) // invoke the MapFn from the gRPC client for a stream of input requests // resultCh -> chan to read responses streamed back // reduceErrCh -> chan for reading any errors encountered during gRPC @@ -94,19 +95,19 @@ func (u *GRPCBasedFlatmap) ApplyMap(ctx context.Context, messageStream []*isb.Re 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 - // 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) - } + //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 - responseCh <- resp + //log.Print("MYDEBUG SENDING TO WRITER ", result.Result.GetUuid()) + responseCh <- result case err := <-reduceErrCh: // We got a context done while processing the gRPC, hence stop processing // The specific case for ctx.Done() is already checked in MapFn @@ -137,18 +138,20 @@ func (u *GRPCBasedFlatmap) ApplyMap(ctx context.Context, messageStream []*isb.Re defer close(flatmapRequests) for _, req := range messageStream { d := u.createMapRequest(req) + //log.Print("MYDEBUG SENDING TO client ", d.GetUuid()) flatmapRequests <- d } }() - return errCh + return doneChan, errCh } -// createMapRequest takes a isb.ReadMessage and returns proto MapRequest -func (u *GRPCBasedFlatmap) createMapRequest(msg *isb.ReadMessage) *flatmappb.MapRequest { +// 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) + //// 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, @@ -156,7 +159,7 @@ func (u *GRPCBasedFlatmap) createMapRequest(msg *isb.ReadMessage) *flatmappb.Map EventTime: timestamppb.New(msg.EventTime), Watermark: timestamppb.New(msg.Watermark), Headers: msg.Headers, - Uuid: uid, + Uuid: reqMsg.Uid, } return d } @@ -168,57 +171,57 @@ func (u *GRPCBasedFlatmap) parseMapResponse(resp *flatmappb.MapResponse) (parsed result := resp.Result eor := result.GetEOR() uid = result.GetUuid() - parentRequest, ok := u.tracker.GetRequest(uid) + //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 { - } + //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), + //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(), - } + //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: taggedMessage, - AckIt: false, - Total: int64(resp.Result.Total), + //ParentMessage: parentRequest, + Uid: uid, + //RespMessage: result, + AckIt: false, + Total: int64(resp.Result.Total), }, false, uid } diff --git a/pkg/flatmap/tracker/tracker.go b/pkg/flatmap/tracker/tracker.go index ce543943e4..0fa26bc748 100644 --- a/pkg/flatmap/tracker/tracker.go +++ b/pkg/flatmap/tracker/tracker.go @@ -6,6 +6,7 @@ import ( "github.com/google/uuid" + "github.com/numaproj/numaflow/pkg/flatmap/types" "github.com/numaproj/numaflow/pkg/isb" ) @@ -14,14 +15,14 @@ type Tracker struct { // TODO(stream): check perf between sync.map and mutex+map responseIdx sync.Map lock sync.RWMutex - m map[string]*isb.ReadMessage + m map[string]*types.RequestFlatmap } func NewTracker() *Tracker { return &Tracker{ //requestMap: sync.Map{}, //responseIdx: sync.Map{}, - m: make(map[string]*isb.ReadMessage), + m: make(map[string]*types.RequestFlatmap), } } @@ -30,15 +31,27 @@ func GetNewId() string { return id.String() } -func (t *Tracker) AddRequest(msg *isb.ReadMessage) 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, + } //t.requestMap.Store(id, msg) - t.Set(id, msg) - return id + t.Set(id, flatmapRequest) + return flatmapRequest } -func (t *Tracker) GetRequest(id string) (*isb.ReadMessage, bool) { +func (t *Tracker) GetRequest(id string) (*types.RequestFlatmap, bool) { return t.Get(id) } @@ -78,14 +91,14 @@ func (t *Tracker) PrintAll() { } } -func (t *Tracker) Get(key string) (*isb.ReadMessage, bool) { +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 *isb.ReadMessage) { +func (t *Tracker) Set(key string, value *types.RequestFlatmap) { t.lock.Lock() defer t.lock.Unlock() t.m[key] = value @@ -96,3 +109,13 @@ func (t *Tracker) Delete(key string) { 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 +} diff --git a/pkg/flatmap/types/types.go b/pkg/flatmap/types/types.go index 1eb744b043..53e885ba1a 100644 --- a/pkg/flatmap/types/types.go +++ b/pkg/flatmap/types/types.go @@ -1,13 +1,17 @@ package types -import "github.com/numaproj/numaflow/pkg/isb" +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 *isb.WriteMessage + RespMessage *flatmappb.MapResponse AckIt bool Total int64 } @@ -21,3 +25,8 @@ type AckMsgFlatmap struct { Message *isb.ReadMessage AckIt bool } + +type RequestFlatmap struct { + Request *isb.ReadMessage + Uid string +} diff --git a/pkg/sdkclient/flatmapper/client.go b/pkg/sdkclient/flatmapper/client.go index 3cc09ba4a1..11c15a6d36 100644 --- a/pkg/sdkclient/flatmapper/client.go +++ b/pkg/sdkclient/flatmapper/client.go @@ -98,6 +98,7 @@ func (c client) MapFn(ctx context.Context, datumStreamCh <-chan *flatmappb.MapRe errCh <- errSDK return } + //log.Print("MYDEBUG GOT FROM grpc ", resp.Result.GetUuid()) responseCh <- resp } } @@ -106,6 +107,7 @@ func (c client) MapFn(ctx context.Context, datumStreamCh <-chan *flatmappb.MapRe // in case there is an error in sending, send it to the error channel for handling go func() { for inputMsg := range datumStreamCh { + //log.Print("MYDEBUG Sending to grpc ", inputMsg.GetUuid()) err := stream.Send(inputMsg) if err != nil { go func(sErr error) { From f39eff04e1400f3df2572c99cc9e14b50b12b462 Mon Sep 17 00:00:00 2001 From: Sidhant Kohli Date: Mon, 10 Jun 2024 16:29:04 -0700 Subject: [PATCH 12/14] read/write isb Signed-off-by: Sidhant Kohli --- pkg/isb/stores/jetstream/reader.go | 12 ++++++++++ pkg/isb/stores/jetstream/writer.go | 35 ++++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/pkg/isb/stores/jetstream/reader.go b/pkg/isb/stores/jetstream/reader.go index 72e643e9cc..7751d4464a 100644 --- a/pkg/isb/stores/jetstream/reader.go +++ b/pkg/isb/stores/jetstream/reader.go @@ -152,6 +152,18 @@ 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} + } + } + } 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..7f54d48d3d 100644 --- a/pkg/isb/stores/jetstream/writer.go +++ b/pkg/isb/stores/jetstream/writer.go @@ -256,10 +256,45 @@ 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) { + 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 { From 424e6e13698d40d9312014a6da2c12b91e479780 Mon Sep 17 00:00:00 2001 From: Sidhant Kohli Date: Tue, 11 Jun 2024 17:24:36 -0700 Subject: [PATCH 13/14] perf changes Signed-off-by: Sidhant Kohli --- pkg/flatmap/flatmap_udf.go | 6 + pkg/flatmap/forward/forward.go | 242 +++++++++++++++++--------- pkg/flatmap/tracker/tracker.go | 11 +- pkg/flatmap/types/types.go | 5 +- pkg/isb/interfaces.go | 1 + pkg/isb/stores/jetstream/reader.go | 1 + pkg/isb/stores/jetstream/writer.go | 35 +++- pkg/isb/stores/redis/write.go | 5 + pkg/isb/stores/simplebuffer/buffer.go | 5 + pkg/sdkclient/flatmapper/client.go | 87 ++++++++- pkg/sdkclient/mapper/client.go | 20 +++ pkg/sinks/blackhole/blackhole.go | 5 + pkg/sinks/kafka/kafka.go | 5 + pkg/sinks/logger/log.go | 5 + pkg/sinks/udsink/sink.go | 5 + pkg/udf/forward/forward.go | 4 +- 16 files changed, 341 insertions(+), 101 deletions(-) diff --git a/pkg/flatmap/flatmap_udf.go b/pkg/flatmap/flatmap_udf.go index 576a2556b4..146f2d6772 100644 --- a/pkg/flatmap/flatmap_udf.go +++ b/pkg/flatmap/flatmap_udf.go @@ -166,8 +166,14 @@ func (u *FlatmapUDFProcessor) Start(ctx context.Context) error { 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 diff --git a/pkg/flatmap/forward/forward.go b/pkg/flatmap/forward/forward.go index e03ebd8e03..810e402c0b 100644 --- a/pkg/flatmap/forward/forward.go +++ b/pkg/flatmap/forward/forward.go @@ -25,6 +25,7 @@ import ( "fmt" "math" "strconv" + "strings" "sync" "sync/atomic" "time" @@ -235,6 +236,8 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { 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, @@ -338,19 +341,23 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { } } + //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) + 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) - //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? + // reset request tracker - isdf.opts.logger.Debugw("forwardAChunk completed", zap.Int("concurrency", isdf.opts.udfConcurrency), zap.Duration("took", time.Since(start))) + 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. @@ -358,28 +365,17 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { // 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() -ackLoop: - for { - select { - case <-ctx.Done(): - break ackLoop - case response, ok := <-ackMsgChan: - if !ok { - break ackLoop - } - //isdf.opts.logger.Info("MYDEBUG: GOT TO ACK ", response.Uid) - ackMessages := []isb.Offset{response.Request.ReadOffset} - if err := isdf.ackFromBuffer(ctx, ackMessages); 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 - } + 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) - //isdf.opts.logger.Info("MYDEBUG: DONE ACK ", response.Uid) - 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()}).Inc() + 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) } } @@ -390,8 +386,6 @@ ackLoop: // 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{}) { - logger := isdf.opts.logger - logger.Info("MYDEBUG: NO WG ACK ROUTINE ", isdf.opts.readBatchSize) doneChan = make(chan struct{}) go func() { defer close(doneChan) @@ -411,56 +405,50 @@ func (isdf *InterStepDataForward) invokeAck(ctx context.Context, ackMsgChan <-ch // 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) { - defer wg.Done() -outerLoop: - for { - select { - case response, ok := <-udfRespCh: - // if the channel is closed, exit from the routine - if !ok { - break outerLoop - } - 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. - // But what if we - isdf.opts.logger.Error("MYDEBUG: NEW ERROR IN WRITE", zap.Error(err)) - } - // 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) { - //isdf.opts.logger.Info("MYDEBUG: SENDING TO ACK ", response.Result.GetTotal()) - ackChan <- trackedRequest - } - 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(1)) + 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 { @@ -513,7 +501,75 @@ func (isdf *InterStepDataForward) forwardToBuffers(ctx context.Context, writeMes return nil } -// ackFromBuffer acknowledges an array of offsets back to fromBufferPartition and is a blocking call or until shutdown has been initiated. +// 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, @@ -569,23 +625,25 @@ 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 } } } - return writeOffsets, nil + 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) { +func (isdf *InterStepDataForward) writeToBuffer(ctx context.Context, toBufferPartition isb.BufferWriter, msg isb.Message) (writeOffsets isb.Offset, err error) { var ( //totalCount int writeCount int @@ -598,14 +656,14 @@ func (isdf *InterStepDataForward) writeToBuffer(ctx context.Context, toBufferPar // EXTRA //var _writeOffsets []isb.Offset = nil //var errs []error = nil - _writeOffsets, errs := toBufferPartition.Write(ctx, []isb.Message{msg}) + _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 = errs[0]; 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 @@ -653,7 +711,7 @@ func (isdf *InterStepDataForward) writeToBuffer(ctx context.Context, toBufferPar if needRetry { isdf.opts.logger.Errorw("Retrying failed messages", - zap.Any("errors", errorArrayToMap(errs)), + zap.Any("errors", errs.Error()), zap.String(metrics.LabelPipeline, isdf.pipelineName), zap.String(metrics.LabelVertex, isdf.vertexName), zap.String(metrics.LabelPartitionName, toBufferPartition.GetName()), @@ -667,8 +725,8 @@ func (isdf *InterStepDataForward) writeToBuffer(ctx context.Context, toBufferPar } } - 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) + //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 } @@ -752,7 +810,7 @@ func (isdf *InterStepDataForward) parseMapResponse(resp *flatmappb.MapResponse, // 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(), isdf.vertexName, result.GetIndex()), + ID: getMessageId(trackedRequest.ReadOffset.String(), isdf.vertexName, result.GetIndex()), Keys: keys, }, Body: isb.Body{ @@ -763,3 +821,13 @@ func (isdf *InterStepDataForward) parseMapResponse(resp *flatmappb.MapResponse, } 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/tracker/tracker.go b/pkg/flatmap/tracker/tracker.go index 0fa26bc748..2512284efe 100644 --- a/pkg/flatmap/tracker/tracker.go +++ b/pkg/flatmap/tracker/tracker.go @@ -43,8 +43,9 @@ 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, + Request: msg, + Uid: id, + ReadOffset: msg.ReadOffset, } //t.requestMap.Store(id, msg) t.Set(id, flatmapRequest) @@ -119,3 +120,9 @@ func (t *Tracker) GetItems() []*isb.ReadMessage { } 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 index 53e885ba1a..7db6bd17b1 100644 --- a/pkg/flatmap/types/types.go +++ b/pkg/flatmap/types/types.go @@ -27,6 +27,7 @@ type AckMsgFlatmap struct { } type RequestFlatmap struct { - Request *isb.ReadMessage - Uid string + 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 7751d4464a..f2db09964c 100644 --- a/pkg/isb/stores/jetstream/reader.go +++ b/pkg/isb/stores/jetstream/reader.go @@ -163,6 +163,7 @@ func (jr *jetStreamReader) Ack(_ context.Context, offsets []isb.Offset) []error return []error{err} } } + return nil } errs := make([]error, len(offsets)) done := make(chan struct{}) diff --git a/pkg/isb/stores/jetstream/writer.go b/pkg/isb/stores/jetstream/writer.go index 7f54d48d3d..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,7 +284,10 @@ 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) { +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 @@ -292,7 +323,7 @@ func (jw *jetStreamWriter) syncWrite(_ context.Context, messages []isb.Message, isbWriteTime.With(metricsLabels).Observe(float64(time.Since(t).Microseconds())) }(time.Now()) if len(messages) == 1 { - routine, err := jw.writeRoutine(messages[0], metricsLabels) + routine, err := jw.WriteRoutine(messages[0], metricsLabels) return []isb.Offset{routine}, []error{err} } var writeOffsets = make([]isb.Offset, len(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/flatmapper/client.go b/pkg/sdkclient/flatmapper/client.go index 11c15a6d36..cf549d3a6d 100644 --- a/pkg/sdkclient/flatmapper/client.go +++ b/pkg/sdkclient/flatmapper/client.go @@ -20,6 +20,8 @@ import ( "context" "errors" "io" + "strconv" + "strings" flatmappb "github.com/numaproj/numaflow-go/pkg/apis/proto/flatmap/v1" "github.com/numaproj/numaflow-go/pkg/info" @@ -53,12 +55,12 @@ func (c client) MapFn(ctx context.Context, datumStreamCh <-chan *flatmappb.MapRe 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 + //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 @@ -107,7 +109,6 @@ func (c client) MapFn(ctx context.Context, datumStreamCh <-chan *flatmappb.MapRe // in case there is an error in sending, send it to the error channel for handling go func() { for inputMsg := range datumStreamCh { - //log.Print("MYDEBUG Sending to grpc ", inputMsg.GetUuid()) err := stream.Send(inputMsg) if err != nil { go func(sErr error) { @@ -127,6 +128,28 @@ func (c client) MapFn(ctx context.Context, datumStreamCh <-chan *flatmappb.MapRe } }() 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 { @@ -160,3 +183,53 @@ func New(serverInfo *info.ServerInfo, inputOptions ...sdkclient.Option) (Client, 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/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..36835896f6 100644 --- a/pkg/udf/forward/forward.go +++ b/pkg/udf/forward/forward.go @@ -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. @@ -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. From 201d5d54eff34992f758ee06777dc1bc4e3b5ca0 Mon Sep 17 00:00:00 2001 From: Sidhant Kohli Date: Mon, 17 Jun 2024 11:25:05 -0700 Subject: [PATCH 14/14] remove wm in sync map Signed-off-by: Sidhant Kohli --- pkg/flatmap/forward/forward.go | 5 +-- pkg/flatmap/rpc/grpc_flatmap.go | 6 +-- pkg/udf/forward/forward.go | 76 ++++++++++++++++----------------- 3 files changed, 43 insertions(+), 44 deletions(-) diff --git a/pkg/flatmap/forward/forward.go b/pkg/flatmap/forward/forward.go index 810e402c0b..e8e1783e9a 100644 --- a/pkg/flatmap/forward/forward.go +++ b/pkg/flatmap/forward/forward.go @@ -333,7 +333,7 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { <-ackDone //TODO(stream): is it fine if we ack the ctrlMessageOffsets in the end? - // ack the control messages, also based on some error do + // 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 { @@ -348,6 +348,7 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { 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))) @@ -355,8 +356,6 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { // reset response tracker isdf.responseTracker = new(sync.Map) - // reset request tracker - isdf.opts.logger.Debugw("forwardAChunk with UDF completed", zap.Int("concurrency", isdf.opts.udfConcurrency), zap.Duration("took", time.Since(start))) } diff --git a/pkg/flatmap/rpc/grpc_flatmap.go b/pkg/flatmap/rpc/grpc_flatmap.go index 1783072c34..28bf8824b4 100644 --- a/pkg/flatmap/rpc/grpc_flatmap.go +++ b/pkg/flatmap/rpc/grpc_flatmap.go @@ -85,8 +85,8 @@ func (u *GRPCBasedFlatmap) ApplyMap(ctx context2.Context, messageStream []*types defer close(doneChan) // invoke the MapFn from the gRPC client for a stream of input requests // resultCh -> chan to read responses streamed back - // reduceErrCh -> chan for reading any errors encountered during gRPC - resultCh, reduceErrCh := u.client.MapFn(ctx, flatmapRequests) + // 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 @@ -108,7 +108,7 @@ func (u *GRPCBasedFlatmap) ApplyMap(ctx context2.Context, messageStream []*types // Forward the received response to the channel //log.Print("MYDEBUG SENDING TO WRITER ", result.Result.GetUuid()) responseCh <- result - case err := <-reduceErrCh: + 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() { diff --git a/pkg/udf/forward/forward.go b/pkg/udf/forward/forward.go index 36835896f6..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) @@ -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 @@ -553,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 }