From d79c0b9a0244c39cbe3106a88924e67b82cd2860 Mon Sep 17 00:00:00 2001 From: Sidhant Kohli Date: Mon, 20 May 2024 17:20:16 -0700 Subject: [PATCH 1/5] init Signed-off-by: Sidhant Kohli --- cmd/commands/processor.go | 30 + pkg/apis/numaflow/v1alpha1/const.go | 4 + pkg/apis/numaflow/v1alpha1/vertex_types.go | 9 + pkg/flatmap/common.go | 171 ++ pkg/flatmap/flatmap_udf.go | 298 ++++ 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 | 693 ++++++++ pkg/flatmap/forward/forward_test.go | 1741 ++++++++++++++++++++ pkg/flatmap/forward/options.go | 98 ++ pkg/flatmap/forward/shutdown.go | 74 + pkg/flatmap/forward/shutdown_test.go | 156 ++ pkg/flatmap/rpc/grpc_flatmap.go | 190 +++ pkg/flatmap/tracker/tracker.go | 61 + pkg/flatmap/types/types.go | 19 + pkg/sdkclient/const.go | 2 + pkg/sdkclient/flatmapper/client.go | 129 ++ pkg/sdkclient/flatmapper/interface.go | 14 + 19 files changed, 3777 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/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..4ad1d6bc4c 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,17 @@ 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{ @@ -111,3 +123,21 @@ func NewProcessorCommand() *cobra.Command { command.Flags().StringVar(&isbSvcType, "isbsvc-type", "", "ISB Service type, e.g. jetstream") return command } + +// +//func reader() <-chan Message { +// out := make(chan Message) +// go func() { +// defer close(out) +// for { +// select { +// <- done { +// return +// } +// out <- readIsb() +// } +// } +// }() +// return out +// +//} diff --git a/pkg/apis/numaflow/v1alpha1/const.go b/pkg/apis/numaflow/v1alpha1/const.go index f129efdf53..bf944e4b8a 100644 --- a/pkg/apis/numaflow/v1alpha1/const.go +++ b/pkg/apis/numaflow/v1alpha1/const.go @@ -201,6 +201,10 @@ 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" PipelineStatusUnknown = "unknown" 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..e4f90d9c28 --- /dev/null +++ b/pkg/flatmap/flatmap_udf.go @@ -0,0 +1,298 @@ +/* +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") + + // TODO(stream) : what all do we need here? + // mapBatchHandler maybe + var ( + readers []isb.BufferReader + writers map[string][]isb.BufferWriter + fromVertexWmStores map[string]store.WatermarkStore + toVertexWmStores map[string]store.WatermarkStore + //mapHandler *rpc.GRPCBasedMap + 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) + } + + // TODO(stream): Maybe this annotation should be for Batching instead + enableFlatmapUdfStream, err := u.VertexInstance.Vertex.FlatmapUdfStreamEnabled() + if err != nil { + return fmt.Errorf("failed to parse Flatmap streaming UDF metadata, %w", err) + } + + maxMessageSize := sharedutil.LookupEnvIntOr(dfv1.EnvGRPCMaxMessageSize, sdkclient.DefaultGRPCMaxMessageSize) + if enableFlatmapUdfStream { + // 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 stream client, %w", err) + } + flatmapHandler = rpc.NewUDSgRPCBasedFlatmap(flatMapClient) + + // Readiness check + if err := flatmapHandler.WaitUntilReady(ctx); err != nil { + return fmt.Errorf("failed on map stream 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{forward.WithLogger(log), + forward.WithUDFStreaming(enableFlatmapUdfStream)} + 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 + // TODO(stream): refactor the interface based on the batch handler etc + df, err := forward.NewInterStepDataForward(u.VertexInstance, readers[index], writers, conditionalForwarder, flatmapHandler, nil, 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 + if enableFlatmapUdfStream { + metricsOpts = metrics.NewMetricsOptions(ctx, u.VertexInstance.Vertex, []metrics.HealthChecker{flatmapHandler}, lagReaders) + } + // TODO(stream): might use for batching + //else { + // metricsOpts = metrics.NewMetricsOptions(ctx, u.VertexInstance.Vertex, []metrics.HealthChecker{mapHandler}, 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..96bfdc2d33 --- /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" +) + +// MapApplier 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 MapApplier interface { + ApplyMap(ctx context.Context, messageStream <-chan *isb.ReadMessage) (<-chan *types.ResponseFlatmap, <-chan error) +} + +// ApplyMapFunc utility function used to create a MapApplier implementation +type ApplyMapFunc func(ctx context.Context, messageStream <-chan *isb.ReadMessage) (<-chan *types.ResponseFlatmap, <-chan error) + +func (f ApplyMapFunc) ApplyMap(ctx context.Context, messageStream <-chan *isb.ReadMessage) (<-chan *types.ResponseFlatmap, <-chan error) { + return f(ctx, messageStream) +} 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..af52c45465 --- /dev/null +++ b/pkg/flatmap/forward/forward.go @@ -0,0 +1,693 @@ +/* +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.MapApplier + mapStreamUDF applier.MapStreamApplier + 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.MapApplier, + applyUDFStream applier.MapStreamApplier, + fetchWatermark fetch.Fetcher, + publishWatermark map[string]publish.Publisher, + idleManager wmb.IdleManager, + opts ...Option) (*InterStepDataForward, error) { + + optsDef := DefaultOptions() + for _, o := range opts { + if err := o(optsDef); 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, + mapStreamUDF: applyUDFStream, + 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: *optsDef, + } + + // Add logger from parent ctx to child context. + isdf.ctx = logging.WithLogger(ctx, optsDef.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 +} + +// readData is the asynchronous reader which would constantly keep reading from the ISB and stream +// the messages on the input channel. +// Channel owned: inputMessages channel +// Error condition: If there is an error while reading, keep trying +// Shutdown flow: On receiving a done signal, we stop the reading process, close the inputMessages channel and return +func (isdf *InterStepDataForward) readData() (<-chan *isb.ReadMessage, chan struct{}) { + log := logging.FromContext(isdf.ctx) + + // inputMessages is the channel on which the data read from the ISB will be written and then + // this channel will be consumed by the producer. + inputMessages := make(chan *isb.ReadMessage, isdf.opts.readBatchSize) + // stopChan is the done channel which is used to stop the processing for the reader goroutine + stopChan := make(chan struct{}) + go func() { + log.Info("MYDEBUG: I'm in read function") + defer close(inputMessages) + defer close(stopChan) + for { + select { + // TODO(stream) : should we check for a context done also here? + //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 + // } + case <-stopChan: + log.Error("MYDEBUG: Stop reading from ISB") + return + default: + ctx := isdf.ctx + // 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() + } + // Keep streaming the messages read to the inputMessages channel + for _, msg := range readMessages { + log.Info("MYDEBUG: streaming in read function", msg.ReadOffset.String()) + inputMessages <- msg + } + } + } + }() + return inputMessages, stopChan +} + +func (isdf *InterStepDataForward) processUdf(inputMessageChan <-chan *isb.ReadMessage) <-chan *types.WriteMsgFlatmap { + ctx := isdf.ctx + // 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. + + logger := isdf.opts.logger + // TODO(stream) : enable idle watermark publishing + + //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 + //} + logger.Info("MYDEBUG: I'm processing here") + + // Send thr requests to the grpc server for results, which are received in the udfRespChan + udfRespChan, err := isdf.flatmapUDF.ApplyMap(ctx, inputMessageChan) + // TODO(stream): check error handling + if err != nil { + + } + writeChan := make(chan *types.WriteMsgFlatmap) + + // create a channel which would be passed to the next buffers for writing, + // errors in messages/no-acks will be propagated from this as well + + go func() { + defer close(writeChan) + for msg := range udfRespChan { + logger.Info("MYDEBUG: Let's send to resp Chan here") + select { + // TODO(stream): add error handling and shutdown here + default: + d := isdf.processWriteMessage(msg, true) + logger.Info("MYDEBUG: Sending to write", string(msg.RespMessage.Payload), "Sending to write", msg.Uid) + writeChan <- d + } + } + }() + return writeChan +} + +// 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 { + // TODO(stream): check shutdown path + 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 +} + +func (isdf *InterStepDataForward) writeAhead(writeMessageCh <-chan *types.WriteMsgFlatmap) <-chan *types.AckMsgFlatmap { + // TODO(stream): check buffered channel size + logger := isdf.opts.logger + ackChan := make(chan *types.AckMsgFlatmap) + go func() { + logger.Info("MYDEBUG: I'm writing to buffer here") + defer close(ackChan) + ctx := isdf.ctx + // TODO: check pnf.forwardResponses + // 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) + var writeOffsets = make(map[string][][]isb.Offset) + + for toVertex := range isdf.toBuffers { + // over allocating to have a predictable pattern + messageToStep[toVertex] = make([][]isb.Message, len(isdf.toBuffers[toVertex])) + writeOffsets[toVertex] = make([][]isb.Offset, len(isdf.toBuffers[toVertex])) + } + // 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 msg := range writeMessageCh { + logger.Info("MYDEBUG: Got a msg on writeMessageCh ", msg.Message.Uid) + if !msg.AckIt { + // TODO(stream): if no ack directly set let's short circuit? + } + parentMessage := msg.Message.ParentMessage + writeMessage := msg.Message.RespMessage + + writeMessage.Headers = parentMessage.Headers + // add vertex name to the ID, since multiple vertices can publish to the same vertex and we need uniqueness across them + // TODO(stream): check if this ID is correct for DEDUP? + writeMessage.ID = fmt.Sprintf("%s-%s-%d", parentMessage.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, parentMessage); err != nil { + // TODO(stream): mark as no ack directly or retry + //return nil, fmt.Errorf("failed at whereToStep, error: %w", err) + } + logger.Info("MYDEBUG: STEP 2 at ACKING ", msg.Message.Uid) + + // Forward the message to the edge buffer (could be multiple edges) + curWriteOffsets, err := isdf.writeToBuffers(ctx, messageToStep) + if err != nil { + // TODO(stream): mark as no ack directly or retry + //return nil, fmt.Errorf("failed to write to toBuffers, error: %w", err) + } + logger.Info("MYDEBUG: STEP 3 at ACKING ", msg.Message.Uid) + // Merge curWriteOffsets into writeOffsets + for vertexName, toVertexBufferOffsets := range curWriteOffsets { + for index, offsets := range toVertexBufferOffsets { + writeOffsets[vertexName][index] = append(writeOffsets[vertexName][index], offsets...) + } + } + logger.Info("MYDEBUG: STEP 4 at ACKING ", msg.Message.Uid) + + // TODO(stream): Message written successfully -> send for ack + ackChan <- &types.AckMsgFlatmap{ + Message: parentMessage, + AckIt: msg.AckIt, + } + logger.Info("MYDEBUG: WRITE sent a msg for ACK ", parentMessage.ReadOffset) + // TODO(stream): publish new watermark + } + }() + return ackChan +} + +func (isdf *InterStepDataForward) ackPrevBuffer(ackMsgChan <-chan *types.AckMsgFlatmap) { + ctx := isdf.ctx + flushTimer := time.NewTicker(isdf.opts.flushDuration) + ackMessages := make([]isb.Offset, 0, isdf.opts.batchSize) + noAckMessages := make([]isb.Offset, 0, isdf.opts.batchSize) + + // should we flush? + var flushAck bool + + // should we flush? + var flushNoAck bool + + // error != nil only when the context is closed, so we can safely return (our write loop will try indefinitely + // unless ctx.Done() happens) +forwardLoop: + for { + select { + case response, ok := <-ackMsgChan: + if !ok { + break forwardLoop + } + + if response.AckIt { + // append the ack message to the array + ackMessages = append(ackMessages, response.Message.ReadOffset) + } else { + // append the ack message to the array + noAckMessages = append(noAckMessages, response.Message.ReadOffset) + + } + + // if the batch size is reached, let's flush + if len(ackMessages) >= isdf.opts.batchSize { + flushAck = true + } + + // if the batch size is reached, let's flush + if len(noAckMessages) >= isdf.opts.batchSize { + flushNoAck = true + } + + case <-flushTimer.C: + // if there are no messages to write, continue + if len(ackMessages) == 0 { + continue + } + + // Since flushTimer is triggered, it is time to flush + flushAck = true + flushNoAck = true + } + + if flushAck { + if err := isdf.ackFromBuffer(ctx, ackMessages); err != nil { + // TODO(stream): we have retried in the ackFromBuffer, should we trigger + // shutdown here then? + } + ackMessages = make([]isb.Offset, 0, isdf.opts.batchSize) + flushAck = false + } + + if flushNoAck { + isdf.noAckMessages(ctx, noAckMessages) + noAckMessages = make([]isb.Offset, 0, isdf.opts.batchSize) + flushNoAck = false + } + } + + // if there are any messages left, forward them to the ISB + if len(ackMessages) > 0 { + if err := isdf.ackFromBuffer(ctx, ackMessages); err != nil { + // TODO(stream): we have retried in the ackFromBuffer, should we trigger + // shutdown here then? + //return + } + ackMessages = make([]isb.Offset, 0, isdf.opts.batchSize) + } + + // if there are any messages left, forward them to the ISB + if len(noAckMessages) > 0 { + isdf.noAckMessages(ctx, noAckMessages) + noAckMessages = make([]isb.Offset, 0, isdf.opts.batchSize) + } + +} + +// 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() + + // TODO(stream): check buffered channel size for each + // reading data + inputMessagesChan, doneReadChan := isdf.readData() + + // processing data + writeMsgChan := isdf.processUdf(inputMessagesChan) + + // writing data to next ISB buffer + ackMsgChan := isdf.writeAhead(writeMsgChan) + // + //// Ack to previous ISB + go isdf.ackPrevBuffer(ackMsgChan) + + // TODO(stream): check ideal way to wait here + <-doneReadChan +} + +// 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] = append(messageToStep[t.ToVertexName][t.ToVertexPartitionIdx], writeMessage.Message) + } + return nil +} + +// 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, messages []isb.Message) (writeOffsets []isb.Offset, err error) { + var ( + totalCount int + writeCount int + writeBytes float64 + ) + totalCount = len(messages) + writeOffsets = make([]isb.Offset, 0, totalCount) + + for { + _writeOffsets, errs := toBufferPartition.Write(ctx, messages) + // 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 { + if err = errs[idx]; 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 = append(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:%d, %v", len(failedMessages), errs) + } + } + } else { + writeCount++ + writeBytes += float64(len(msg.Payload)) + // we support write offsets only for jetstream + if _writeOffsets != nil { + writeOffsets = append(writeOffsets, _writeOffsets[idx]) + } + } + } + + 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 + messages = 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 +} + +// 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 +} + +// 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 +} + +// noAckMessages no-acks all the read offsets of failed messages. +func (isdf *InterStepDataForward) noAckMessages(ctx context.Context, failedMessages []isb.Offset) { + isdf.fromBufferPartition.NoAck(ctx, failedMessages) +} + +func (isdf *InterStepDataForward) processWriteMessage(msg *types.ResponseFlatmap, ackIt bool) *types.WriteMsgFlatmap { + + return &types.WriteMsgFlatmap{ + Message: msg, + AckIt: ackIt, + } + +} diff --git a/pkg/flatmap/forward/forward_test.go b/pkg/flatmap/forward/forward_test.go new file mode 100644 index 0000000000..c9b0c35d46 --- /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{}, 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{}, &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{}, 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{}, 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{}, 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{}, 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{}, 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{}, 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{}, 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{}, 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{}, 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{}, 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..d078127991 --- /dev/null +++ b/pkg/flatmap/forward/options.go @@ -0,0 +1,98 @@ +/* +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 + // TODO(stream): Should be used for batch? + // enableMapUdfStream indicates whether the message streaming is enabled or not for map UDF processing + enableMapUdfStream bool + batchSize int + flushDuration time.Duration +} + +type Option func(*options) error + +func DefaultOptions() *options { + return &options{ + readBatchSize: dfv1.DefaultReadBatchSize, + udfConcurrency: dfv1.DefaultReadBatchSize, + retryInterval: time.Millisecond, + logger: logging.NewLogger(), + enableMapUdfStream: false, + // TODO(stream): check the optimization for these vals + batchSize: dfv1.DefaultReadBatchSize, + flushDuration: dfv1.DefaultPnfFlushDuration, + } +} + +// 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..9ea2f923aa --- /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{}, 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{}, 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/grpc_flatmap.go b/pkg/flatmap/rpc/grpc_flatmap.go new file mode 100644 index 0000000000..ac3e21119b --- /dev/null +++ b/pkg/flatmap/rpc/grpc_flatmap.go @@ -0,0 +1,190 @@ +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" + "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 +} + +func NewUDSgRPCBasedFlatmap(client flatmapper.Client) *GRPCBasedFlatmap { + return &GRPCBasedFlatmap{client: client, tracker: tracker.NewTracker()} +} + +// 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 { + log := 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 { + log.Infof("waiting for reduce udf to be ready: %v", err) + time.Sleep(1 * time.Second) + } + } + } +} + +func (u *GRPCBasedFlatmap) ApplyMap(ctx context.Context, messageStream <-chan *isb.ReadMessage) (<-chan *types.ResponseFlatmap, <-chan error) { + var ( + errCh = make(chan error) + responseCh = make(chan *types.ResponseFlatmap) + mapRequestCh = make(chan *flatmappb.MapRequest) + ) + + // invoke the MapFn method with mapRequestCh channel and send the result to responseCh channel + // and any error to errCh channel + go func() { + log.Println("MYDEBUG: I'm processing here") + index := 0 + resultCh, mapErrCh := u.client.MapFn(ctx, mapRequestCh) + for { + select { + case result, ok := <-resultCh: + // TODO(stream): Check error handling here + if !ok || result == nil { + errCh = nil + // if the resultCh channel is closed, close the responseCh and return + close(responseCh) + return + } + // create a unique message id for each response message which will be used for deduplication + index++ + responseCh <- u.parseMapResponse(result) + // TODO(stream): We need to remove the request message from the tracker once this is completed. + // As we are streaming messages, we need to have some control field to indicate that this is completed + // now, we can do that in the SDK itself. + + case err := <-mapErrCh: + // TODO(stream): Check error handling here + // 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) + //} + } + } + }() + + // create ReduceRequest from TimedWindowRequest and send it to reduceRequests channel for AsyncReduceFn + go func() { + // after reading all the messages from the requestsStream or if ctx was canceled close the reduceRequests channel + defer func() { + close(mapRequestCh) + }() + for { + select { + case msg, ok := <-messageStream: + log.Println("MYDEBUG: reading for messages here") + // if the requestsStream is closed or if the message is nil, return + if !ok || msg == nil { + //return + } + + d := u.createFlatmapRequest(msg) + // send the datum to reduceRequests channel, handle the case when the context is canceled + select { + // TODO(stream): Check the context end here + case mapRequestCh <- d: + log.Println("MYDEBUG: send the message here", d.Uuid) + //case <-ctx.Done(): + // return + } + // TODO(stream): Check the context end here, need to invoke shutdown + //case <-ctx.Done(): // if the context is done, don't send any more datum to reduceRequests channel + // return + } + } + }() + + return responseCh, errCh + +} + +func (u *GRPCBasedFlatmap) createFlatmapRequest(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) *types.ResponseFlatmap { + result := resp.Result + uid := result.GetUuid() + parentRequest, ok := u.tracker.GetRequest(uid) + // TODO(stream): check what should be path for !ok + if !ok { + + } + idx, present := u.tracker.GetIdx(uid) + if !present { + u.tracker.NewResponse(uid) + idx = 1 + } + 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-%d", parentRequest.ReadOffset.String(), idx), + Keys: keys, + }, + Body: isb.Body{ + Payload: result.GetValue(), + }, + }, + Tags: result.GetTags(), + } + return &types.ResponseFlatmap{ + ParentMessage: parentRequest, + Uid: uid, + RespMessage: taggedMessage, + } +} diff --git a/pkg/flatmap/tracker/tracker.go b/pkg/flatmap/tracker/tracker.go new file mode 100644 index 0000000000..d8dbed2785 --- /dev/null +++ b/pkg/flatmap/tracker/tracker.go @@ -0,0 +1,61 @@ +package tracker + +import ( + "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 +} + +func NewTracker() *Tracker { + return &Tracker{ + requestMap: sync.Map{}, + responseIdx: sync.Map{}, + } +} + +func GetNewId() string { + id, _ := uuid.NewUUID() + return id.String() +} + +func (t *Tracker) AddRequest(msg *isb.ReadMessage) string { + id := GetNewId() + t.requestMap.Store(id, msg) + return id +} + +func (t *Tracker) GetRequest(id string) (*isb.ReadMessage, bool) { + val, ok := t.requestMap.Load(id) + return val.(*isb.ReadMessage), ok +} + +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 + +} diff --git a/pkg/flatmap/types/types.go b/pkg/flatmap/types/types.go new file mode 100644 index 0000000000..6453ffc30e --- /dev/null +++ b/pkg/flatmap/types/types.go @@ -0,0 +1,19 @@ +package types + +import "github.com/numaproj/numaflow/pkg/isb" + +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..e74ff54489 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/flatmap-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..83ab859faa --- /dev/null +++ b/pkg/sdkclient/flatmapper/client.go @@ -0,0 +1,129 @@ +package flatmapper + +import ( + "errors" + "io" + + flatmappb "github.com/numaproj/numaflow-go/pkg/apis/proto/flatmap/v1" + "github.com/numaproj/numaflow-go/pkg/info" + "golang.org/x/net/context" + "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) 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 +} +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) + ) + + // stream the messages to server + stream, err := c.grpcClt.MapFn(ctx) + + if err != nil { + go func(sErr error) { + errCh <- sdkerr.ToUDFErr("c.grpcClt.MapFn", sErr) + }(err) + } + + // read from the datumStreamCh channel and send it to the server stream + go func() { + var sendErr error + outerLoop: + for { + select { + case <-ctx.Done(): + return + case datum, ok := <-datumStreamCh: + if !ok { + break outerLoop + } + // TODO: figure out why send is getting EOF (could be because the client has already handled SIGTERM) + if sendErr = stream.Send(datum); sendErr != nil && !errors.Is(sendErr, io.EOF) { + errCh <- sdkerr.ToUDFErr("MapFn stream.Send()", sendErr) + return + } + } + } + + select { + case <-ctx.Done(): + return + default: + sendErr = stream.CloseSend() + if sendErr != nil && !errors.Is(sendErr, io.EOF) { + errCh <- sdkerr.ToUDFErr("MapFn stream.CloseSend()", sendErr) + } + } + }() + + // read the response from the server stream and send it to responseCh channel + // any error is sent to errCh channel + go func() { + var resp *flatmappb.MapResponse + var recvErr error + for { + select { + case <-ctx.Done(): + errCh <- ctx.Err() + return + default: + resp, recvErr = stream.Recv() + // if the stream is closed, close the responseCh return + if errors.Is(recvErr, io.EOF) { + // nil channel will never be selected + errCh = nil + close(responseCh) + return + } + if recvErr != nil { + errCh <- sdkerr.ToUDFErr("MapFn stream.Recv()", recvErr) + } + responseCh <- resp + } + } + }() + + return responseCh, errCh +} + +// 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 a78966e4a40966d8aa4ba36a6529962b60f532c2 Mon Sep 17 00:00:00 2001 From: Sidhant Kohli Date: Tue, 21 May 2024 17:50:56 -0700 Subject: [PATCH 2/5] init Signed-off-by: Sidhant Kohli --- pkg/flatmap/flatmap_udf.go | 3 +- pkg/flatmap/forward/forward.go | 659 +++++++++++++++++++++++++------- pkg/flatmap/rpc/grpc_flatmap.go | 26 +- pkg/flatmap/tracker/tracker.go | 4 + 4 files changed, 553 insertions(+), 139 deletions(-) diff --git a/pkg/flatmap/flatmap_udf.go b/pkg/flatmap/flatmap_udf.go index e4f90d9c28..6046b78961 100644 --- a/pkg/flatmap/flatmap_udf.go +++ b/pkg/flatmap/flatmap_udf.go @@ -215,9 +215,10 @@ func (u *FlatMapUDFProcessor) Start(ctx context.Context) error { opts = append(opts, forward.WithUDFConcurrency(int(*x.ReadBatchSize))) } } + // create a forwarder for each partition // TODO(stream): refactor the interface based on the batch handler etc - df, err := forward.NewInterStepDataForward(u.VertexInstance, readers[index], writers, conditionalForwarder, flatmapHandler, nil, fetchWatermark, publishWatermark, idleManager, opts...) + df, err := forward.NewInterStepDataForward(u.VertexInstance, readers[index], writers, conditionalForwarder, flatmapHandler, nil, fetchWatermark, publishWatermark, idleManager, conditionalForwarder, opts...) if err != nil { return err } diff --git a/pkg/flatmap/forward/forward.go b/pkg/flatmap/forward/forward.go index af52c45465..5f88e0472d 100644 --- a/pkg/flatmap/forward/forward.go +++ b/pkg/flatmap/forward/forward.go @@ -29,6 +29,7 @@ import ( "time" "go.uber.org/zap" + "golang.org/x/sync/errgroup" "k8s.io/apimachinery/pkg/util/wait" dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1" @@ -66,7 +67,8 @@ 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 + wmbChecker wmb.WMBChecker + whereToDecider forwarder.ToWhichStepDecider Shutdown } @@ -81,6 +83,7 @@ func NewInterStepDataForward( fetchWatermark fetch.Fetcher, publishWatermark map[string]publish.Publisher, idleManager wmb.IdleManager, + whereToDecider forwarder.ToWhichStepDecider, opts ...Option) (*InterStepDataForward, error) { optsDef := DefaultOptions() @@ -112,7 +115,8 @@ func NewInterStepDataForward( Shutdown: Shutdown{ rwlock: new(sync.RWMutex), }, - opts: *optsDef, + whereToDecider: whereToDecider, + opts: *optsDef, } // Add logger from parent ctx to child context. @@ -131,7 +135,7 @@ func NewInterStepDataForward( // Error condition: If there is an error while reading, keep trying // Shutdown flow: On receiving a done signal, we stop the reading process, close the inputMessages channel and return func (isdf *InterStepDataForward) readData() (<-chan *isb.ReadMessage, chan struct{}) { - log := logging.FromContext(isdf.ctx) + log := isdf.opts.logger // inputMessages is the channel on which the data read from the ISB will be written and then // this channel will be consumed by the producer. @@ -139,7 +143,7 @@ func (isdf *InterStepDataForward) readData() (<-chan *isb.ReadMessage, chan stru // stopChan is the done channel which is used to stop the processing for the reader goroutine stopChan := make(chan struct{}) go func() { - log.Info("MYDEBUG: I'm in read function") + //log.Info("MYDEBUG: I'm in read function") defer close(inputMessages) defer close(stopChan) for { @@ -156,7 +160,7 @@ func (isdf *InterStepDataForward) readData() (<-chan *isb.ReadMessage, chan stru // return // } case <-stopChan: - log.Error("MYDEBUG: Stop reading from ISB") + //log.Error("MYDEBUG: Stop reading from ISB") return default: ctx := isdf.ctx @@ -171,7 +175,7 @@ func (isdf *InterStepDataForward) readData() (<-chan *isb.ReadMessage, chan stru } // Keep streaming the messages read to the inputMessages channel for _, msg := range readMessages { - log.Info("MYDEBUG: streaming in read function", msg.ReadOffset.String()) + log.Info("MYDEBUG: streaming in read function ", msg.ReadOffset.String(), " ", time.Now().UnixNano()) inputMessages <- msg } } @@ -216,7 +220,7 @@ func (isdf *InterStepDataForward) processUdf(inputMessageChan <-chan *isb.ReadMe // } // return //} - logger.Info("MYDEBUG: I'm processing here") + //logger.Info("MYDEBUG: I'm processing here") // Send thr requests to the grpc server for results, which are received in the udfRespChan udfRespChan, err := isdf.flatmapUDF.ApplyMap(ctx, inputMessageChan) @@ -232,12 +236,13 @@ func (isdf *InterStepDataForward) processUdf(inputMessageChan <-chan *isb.ReadMe go func() { defer close(writeChan) for msg := range udfRespChan { - logger.Info("MYDEBUG: Let's send to resp Chan here") + //logger.Info("MYDEBUG: Let's send to resp Chan here") select { // TODO(stream): add error handling and shutdown here default: d := isdf.processWriteMessage(msg, true) - logger.Info("MYDEBUG: Sending to write", string(msg.RespMessage.Payload), "Sending to write", msg.Uid) + logger.Info("MYDEBUG: Sending to write ", msg.ParentMessage.ReadOffset.String(), " ", time.Now().UnixNano()) + //logger.Info("MYDEBUG: Sending to write", string(msg.RespMessage.Payload), "Sending to write", msg.Uid) writeChan <- d } } @@ -308,163 +313,377 @@ type readWriteMessagePair struct { udfError error } -func (isdf *InterStepDataForward) writeAhead(writeMessageCh <-chan *types.WriteMsgFlatmap) <-chan *types.AckMsgFlatmap { - // TODO(stream): check buffered channel size - logger := isdf.opts.logger - ackChan := make(chan *types.AckMsgFlatmap) - go func() { - logger.Info("MYDEBUG: I'm writing to buffer here") - defer close(ackChan) - ctx := isdf.ctx - // TODO: check pnf.forwardResponses - // 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) - var writeOffsets = make(map[string][][]isb.Offset) - - for toVertex := range isdf.toBuffers { - // over allocating to have a predictable pattern - messageToStep[toVertex] = make([][]isb.Message, len(isdf.toBuffers[toVertex])) - writeOffsets[toVertex] = make([][]isb.Offset, len(isdf.toBuffers[toVertex])) - } - // 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 msg := range writeMessageCh { - logger.Info("MYDEBUG: Got a msg on writeMessageCh ", msg.Message.Uid) - if !msg.AckIt { - // TODO(stream): if no ack directly set let's short circuit? - } - parentMessage := msg.Message.ParentMessage - writeMessage := msg.Message.RespMessage - - writeMessage.Headers = parentMessage.Headers - // add vertex name to the ID, since multiple vertices can publish to the same vertex and we need uniqueness across them - // TODO(stream): check if this ID is correct for DEDUP? - writeMessage.ID = fmt.Sprintf("%s-%s-%d", parentMessage.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, parentMessage); err != nil { - // TODO(stream): mark as no ack directly or retry - //return nil, fmt.Errorf("failed at whereToStep, error: %w", err) - } - logger.Info("MYDEBUG: STEP 2 at ACKING ", msg.Message.Uid) - - // Forward the message to the edge buffer (could be multiple edges) - curWriteOffsets, err := isdf.writeToBuffers(ctx, messageToStep) - if err != nil { - // TODO(stream): mark as no ack directly or retry - //return nil, fmt.Errorf("failed to write to toBuffers, error: %w", err) - } - logger.Info("MYDEBUG: STEP 3 at ACKING ", msg.Message.Uid) - // Merge curWriteOffsets into writeOffsets - for vertexName, toVertexBufferOffsets := range curWriteOffsets { - for index, offsets := range toVertexBufferOffsets { - writeOffsets[vertexName][index] = append(writeOffsets[vertexName][index], offsets...) - } - } - logger.Info("MYDEBUG: STEP 4 at ACKING ", msg.Message.Uid) - - // TODO(stream): Message written successfully -> send for ack - ackChan <- &types.AckMsgFlatmap{ - Message: parentMessage, - AckIt: msg.AckIt, - } - logger.Info("MYDEBUG: WRITE sent a msg for ACK ", parentMessage.ReadOffset) - // TODO(stream): publish new watermark - } - }() - return ackChan -} - -func (isdf *InterStepDataForward) ackPrevBuffer(ackMsgChan <-chan *types.AckMsgFlatmap) { - ctx := isdf.ctx +func (isdf *InterStepDataForward) writeRoutine(ctx context.Context, writeMessageCh <-chan *types.WriteMsgFlatmap, ackChan chan<- *types.AckMsgFlatmap) { flushTimer := time.NewTicker(isdf.opts.flushDuration) - ackMessages := make([]isb.Offset, 0, isdf.opts.batchSize) - noAckMessages := make([]isb.Offset, 0, isdf.opts.batchSize) - - // should we flush? - var flushAck bool + writeMessages := make([]*isb.WriteMessage, 0, isdf.opts.batchSize) + readOffsets := make([]*isb.ReadMessage, 0, isdf.opts.batchSize) + logger := isdf.opts.logger // should we flush? - var flushNoAck bool + var flush bool + //logger.Info("MYDEBUG: I'm writing to buffer here") // error != nil only when the context is closed, so we can safely return (our write loop will try indefinitely // unless ctx.Done() happens) + flush = true forwardLoop: for { select { - case response, ok := <-ackMsgChan: + case response, ok := <-writeMessageCh: if !ok { + // TODO(stream): check the error logic here break forwardLoop } - - if response.AckIt { - // append the ack message to the array - ackMessages = append(ackMessages, response.Message.ReadOffset) - } else { - // append the ack message to the array - noAckMessages = append(noAckMessages, response.Message.ReadOffset) - + if !response.AckIt { + // TODO(stream): if no ack directly set let's short circuit? } - // if the batch size is reached, let's flush - if len(ackMessages) >= isdf.opts.batchSize { - flushAck = true - } + // append the write message to the array + writeMessages = append(writeMessages, response.Message.RespMessage) + readOffsets = append(readOffsets, response.Message.ParentMessage) // if the batch size is reached, let's flush - if len(noAckMessages) >= isdf.opts.batchSize { - flushNoAck = true + if len(writeMessages) >= isdf.opts.batchSize { + flush = true } case <-flushTimer.C: // if there are no messages to write, continue - if len(ackMessages) == 0 { + if len(writeMessages) == 0 { continue } // Since flushTimer is triggered, it is time to flush - flushAck = true - flushNoAck = true + flush = true } - if flushAck { - if err := isdf.ackFromBuffer(ctx, ackMessages); err != nil { - // TODO(stream): we have retried in the ackFromBuffer, should we trigger - // shutdown here then? + if flush { + if err := isdf.forwardToBuffers(ctx, &writeMessages); err != nil { + // TODO(stream): mark as no ack directly or retry } - ackMessages = make([]isb.Offset, 0, isdf.opts.batchSize) - flushAck = false - } - if flushNoAck { - isdf.noAckMessages(ctx, noAckMessages) - noAckMessages = make([]isb.Offset, 0, isdf.opts.batchSize) - flushNoAck = false + for _, offset := range readOffsets { + logger.Info("MYDEBUG: Sending to ack ", offset.ReadOffset.String(), " ", time.Now().UnixNano()) + //logger.Info("MYDEBUG: STEP 1 at ACKING ", offset.ReadOffset) + ackChan <- &types.AckMsgFlatmap{ + Message: offset, + // TODO(stream): should we write for no-ack? + AckIt: true, + } + } + // clear the writeMessages + writeMessages = make([]*isb.WriteMessage, 0, isdf.opts.batchSize) + readOffsets = make([]*isb.ReadMessage, 0, isdf.opts.batchSize) + //flush = false } } // if there are any messages left, forward them to the ISB - if len(ackMessages) > 0 { - if err := isdf.ackFromBuffer(ctx, ackMessages); err != nil { - // TODO(stream): we have retried in the ackFromBuffer, should we trigger - // shutdown here then? - //return + if len(writeMessages) > 0 { + if err := isdf.forwardToBuffers(ctx, &writeMessages); err != nil { + // TODO(stream): mark as no ack directly or retry } - ackMessages = make([]isb.Offset, 0, isdf.opts.batchSize) + for _, offset := range readOffsets { + //logger.Info("MYDEBUG: STEP 1 at ACKING ", offset.ReadOffset) + ackChan <- &types.AckMsgFlatmap{ + Message: offset, + // TODO(stream): should we write for no-ack? + AckIt: true, + } + } + // clear the writeMessages + writeMessages = make([]*isb.WriteMessage, 0, isdf.opts.batchSize) + readOffsets = make([]*isb.ReadMessage, 0, isdf.opts.batchSize) + //flush = false } +} + +func (isdf *InterStepDataForward) writeAhead(writeMessageCh <-chan *types.WriteMsgFlatmap) <-chan *types.AckMsgFlatmap { + // TODO(stream): check buffered channel size + //logger := isdf.opts.logger + ackChan := make(chan *types.AckMsgFlatmap) + ctx := isdf.ctx + + go func() { + defer close(ackChan) + group := sync.WaitGroup{} + for i := 0; i < isdf.opts.batchSize; i++ { + group.Add(1) + go isdf.writeRoutine(ctx, writeMessageCh, ackChan) + } + group.Wait() + }() + return ackChan + + //go func() { + // defer close(ackChan) + // flushTimer := time.NewTicker(isdf.opts.flushDuration) + // writeMessages := make([]*isb.WriteMessage, 0, isdf.opts.batchSize) + // readOffsets := make([]*isb.ReadMessage, 0, isdf.opts.batchSize) + // + // // should we flush? + // var flush bool + // //logger.Info("MYDEBUG: I'm writing to buffer here") + // + // // error != nil only when the context is closed, so we can safely return (our write loop will try indefinitely + // // unless ctx.Done() happens) + // flush = true + //forwardLoop: + // for { + // select { + // case response, ok := <-writeMessageCh: + // if !ok { + // // TODO(stream): check the error logic here + // break forwardLoop + // } + // if !response.AckIt { + // // TODO(stream): if no ack directly set let's short circuit? + // } + // + // // append the write message to the array + // writeMessages = append(writeMessages, response.Message.RespMessage) + // readOffsets = append(readOffsets, response.Message.ParentMessage) + // + // // if the batch size is reached, let's flush + // if len(writeMessages) >= isdf.opts.batchSize { + // flush = true + // } + // + // case <-flushTimer.C: + // // if there are no messages to write, continue + // if len(writeMessages) == 0 { + // continue + // } + // + // // Since flushTimer is triggered, it is time to flush + // flush = true + // } + // + // if flush { + // if err := isdf.forwardToBuffers(ctx, &writeMessages); err != nil { + // // TODO(stream): mark as no ack directly or retry + // } + // + // for _, offset := range readOffsets { + // logger.Info("MYDEBUG: Sending to ack ", offset.ReadOffset.String(), " ", time.Now().UnixNano()) + // //logger.Info("MYDEBUG: STEP 1 at ACKING ", offset.ReadOffset) + // ackChan <- &types.AckMsgFlatmap{ + // Message: offset, + // // TODO(stream): should we write for no-ack? + // AckIt: true, + // } + // } + // // clear the writeMessages + // writeMessages = make([]*isb.WriteMessage, 0, isdf.opts.batchSize) + // readOffsets = make([]*isb.ReadMessage, 0, isdf.opts.batchSize) + // //flush = false + // } + // } + // + // // if there are any messages left, forward them to the ISB + // if len(writeMessages) > 0 { + // if err := isdf.forwardToBuffers(ctx, &writeMessages); err != nil { + // // TODO(stream): mark as no ack directly or retry + // } + // for _, offset := range readOffsets { + // //logger.Info("MYDEBUG: STEP 1 at ACKING ", offset.ReadOffset) + // ackChan <- &types.AckMsgFlatmap{ + // Message: offset, + // // TODO(stream): should we write for no-ack? + // AckIt: true, + // } + // } + // // clear the writeMessages + // writeMessages = make([]*isb.WriteMessage, 0, isdf.opts.batchSize) + // readOffsets = make([]*isb.ReadMessage, 0, isdf.opts.batchSize) + // //flush = false + // } + //}() + //go func() { + // logger.Info("MYDEBUG: I'm writing to buffer here") + // defer close(ackChan) + // ctx := isdf.ctx + // // TODO: check pnf.forwardResponses + // // 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) + // var writeOffsets = make(map[string][][]isb.Offset) + // + // for toVertex := range isdf.toBuffers { + // // over allocating to have a predictable pattern + // messageToStep[toVertex] = make([][]isb.Message, len(isdf.toBuffers[toVertex])) + // writeOffsets[toVertex] = make([][]isb.Offset, len(isdf.toBuffers[toVertex])) + // } + // // 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 msg := range writeMessageCh { + // logger.Info("MYDEBUG: Got a msg on writeMessageCh ", msg.Message.Uid) + // if !msg.AckIt { + // // TODO(stream): if no ack directly set let's short circuit? + // } + // parentMessage := msg.Message.ParentMessage + // writeMessage := msg.Message.RespMessage + // + // writeMessage.Headers = parentMessage.Headers + // // add vertex name to the ID, since multiple vertices can publish to the same vertex and we need uniqueness across them + // // TODO(stream): check if this ID is correct for DEDUP? + // logger.Info("MYDEBUG: DEDUP ID ", writeMessage.ID) + // //writeMessage.ID = fmt.Sprintf("%s-%s-%d", parentMessage.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, parentMessage); err != nil { + // // TODO(stream): mark as no ack directly or retry + // //return nil, fmt.Errorf("failed at whereToStep, error: %w", err) + // } + // logger.Info("MYDEBUG: STEP 2 at ACKING ", msg.Message.Uid) + // + // // Forward the message to the edge buffer (could be multiple edges) + // curWriteOffsets, err := isdf.writeToBuffers(ctx, messageToStep) + // if err != nil { + // // TODO(stream): mark as no ack directly or retry + // //return nil, fmt.Errorf("failed to write to toBuffers, error: %w", err) + // } + // logger.Info("MYDEBUG: STEP 3 at ACKING ", msg.Message.Uid) + // // Merge curWriteOffsets into writeOffsets + // for vertexName, toVertexBufferOffsets := range curWriteOffsets { + // for index, offsets := range toVertexBufferOffsets { + // writeOffsets[vertexName][index] = append(writeOffsets[vertexName][index], offsets...) + // } + // } + // logger.Info("MYDEBUG: STEP 4 at ACKING ", len(curWriteOffsets)) + // + // // TODO(stream): Message written successfully -> send for ack + // ackChan <- &types.AckMsgFlatmap{ + // Message: parentMessage, + // AckIt: msg.AckIt, + // } + // logger.Info("MYDEBUG: WRITE sent a msg for ACK ", parentMessage.ReadOffset) + // // TODO(stream): publish new watermark + // } + //}() + //return ackChan + +} + +func (isdf *InterStepDataForward) ackRoutine(ctx context.Context, ackMsgChan <-chan *types.AckMsgFlatmap) { + logger := isdf.opts.logger +forwardLoop: + for { + select { + case response, ok := <-ackMsgChan: + if !ok { + break forwardLoop + } + if response.AckIt { + ackMessages := []isb.Offset{response.Message.ReadOffset} + if err := isdf.ackFromBuffer(ctx, ackMessages); err != nil { + // TODO(stream): we have retried in the ackFromBuffer, should we trigger + // shutdown here then? + } + logger.Info("MYDEBUG: Done with ack ", ackMessages, " ", time.Now().UnixNano()) + } else { + noAckMessages := []isb.Offset{response.Message.ReadOffset} + isdf.noAckMessages(ctx, noAckMessages) + } + } - // if there are any messages left, forward them to the ISB - if len(noAckMessages) > 0 { - isdf.noAckMessages(ctx, noAckMessages) - noAckMessages = make([]isb.Offset, 0, isdf.opts.batchSize) } } +//func (isdf *InterStepDataForward) ackPrevBuffer(ackMsgChan <-chan *types.AckMsgFlatmap) { +// ctx := isdf.ctx +// flushTimer := time.NewTicker(isdf.opts.flushDuration) +// ackMessages := make([]isb.Offset, 0, isdf.opts.batchSize) +// noAckMessages := make([]isb.Offset, 0, isdf.opts.batchSize) +// logger := isdf.opts.logger +// +// // should we flush? +// var flushAck bool +// +// // should we flush? +// var flushNoAck bool +// +// flushAck = true +// flushNoAck = true +// // error != nil only when the context is closed, so we can safely return (our write loop will try indefinitely +// // unless ctx.Done() happens) +//forwardLoop: +// for { +// select { +// case response, ok := <-ackMsgChan: +// if !ok { +// break forwardLoop +// } +// +// if response.AckIt { +// // append the ack message to the array +// ackMessages = append(ackMessages, response.Message.ReadOffset) +// } else { +// // append the ack message to the array +// noAckMessages = append(noAckMessages, response.Message.ReadOffset) +// +// } +// +// // if the batch size is reached, let's flush +// if len(ackMessages) >= isdf.opts.batchSize { +// flushAck = true +// } +// +// // if the batch size is reached, let's flush +// if len(noAckMessages) >= isdf.opts.batchSize { +// flushNoAck = true +// } +// +// case <-flushTimer.C: +// // if there are no messages to write, continue +// if len(ackMessages) == 0 { +// continue +// } +// +// // Since flushTimer is triggered, it is time to flush +// flushAck = true +// flushNoAck = true +// } +// +// if flushAck { +// if err := isdf.ackFromBuffer(ctx, ackMessages); err != nil { +// // TODO(stream): we have retried in the ackFromBuffer, should we trigger +// // shutdown here then? +// } +// logger.Info("MYDEBUG: Sending to ack ", ackMessages, " ", time.Now().UnixNano()) +// ackMessages = make([]isb.Offset, 0, isdf.opts.batchSize) +// //flushAck = false +// } +// +// if flushNoAck { +// isdf.noAckMessages(ctx, noAckMessages) +// noAckMessages = make([]isb.Offset, 0, isdf.opts.batchSize) +// //flushNoAck = false +// } +// } +// +// // if there are any messages left, forward them to the ISB +// if len(ackMessages) > 0 { +// if err := isdf.ackFromBuffer(ctx, ackMessages); err != nil { +// // TODO(stream): we have retried in the ackFromBuffer, should we trigger +// // shutdown here then? +// //return +// } +// ackMessages = make([]isb.Offset, 0, isdf.opts.batchSize) +// } +// +// // if there are any messages left, forward them to the ISB +// if len(noAckMessages) > 0 { +// isdf.noAckMessages(ctx, noAckMessages) +// noAckMessages = make([]isb.Offset, 0, isdf.opts.batchSize) +// } + +//} + // 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, @@ -481,9 +700,16 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { // writing data to next ISB buffer ackMsgChan := isdf.writeAhead(writeMsgChan) - // - //// Ack to previous ISB - go isdf.ackPrevBuffer(ackMsgChan) + + // Ack to previous ISB + go func() { + group := sync.WaitGroup{} + for i := 0; i < isdf.opts.batchSize; i++ { + group.Add(1) + go isdf.ackRoutine(isdf.ctx, ackMsgChan) + } + group.Wait() + }() // TODO(stream): check ideal way to wait here <-doneReadChan @@ -684,10 +910,181 @@ func (isdf *InterStepDataForward) noAckMessages(ctx context.Context, failedMessa } func (isdf *InterStepDataForward) processWriteMessage(msg *types.ResponseFlatmap, ackIt bool) *types.WriteMsgFlatmap { - return &types.WriteMsgFlatmap{ Message: msg, AckIt: ackIt, } +} + +// forwardToBuffers writes the messages to the ISBs concurrently for each partition. +func (isdf *InterStepDataForward) forwardToBuffers(ctx context.Context, writeMessages *[]*isb.WriteMessage) error { + if len(*writeMessages) == 0 { + return nil + } + messagesToStep := isdf.whereToStepNew(*writeMessages) + // parallel writes to each ISB + var mu sync.Mutex + // use error group + var eg errgroup.Group + for key, values := range messagesToStep { + for index, messages := range values { + if len(messages) == 0 { + continue + } + + func(toVertexName string, toVertexPartitionIdx int32, resultMessages []isb.Message) { + eg.Go(func() error { + _, err := isdf.writeToBufferNew(ctx, toVertexName, toVertexPartitionIdx, resultMessages) + if err != nil { + return err + } + mu.Lock() + // TODO: do we need lock? isn't each buffer isolated since we do sequential per ISB? + //isdf.latestWriteOffsets[toVertexName][toVertexPartitionIdx] = offsets + mu.Unlock() + return nil + }) + }(key, int32(index), messages) + } + } + + // wait until all the writer go routines return + if err := eg.Wait(); err != nil { + return err + } + + return nil +} + +// whereToStep assigns a message to the ISBs based on the Message.Keys. +func (isdf *InterStepDataForward) whereToStepNew(writeMessages []*isb.WriteMessage) map[string][][]isb.Message { + // writer doesn't accept array of pointers + messagesToStep := make(map[string][][]isb.Message) + + var to []forwarder.VertexBuffer + var err error + for _, msg := range writeMessages { + to, err = isdf.whereToDecider.WhereTo(msg.Keys, msg.Tags, msg.ID) + if err != nil { + metrics.PlatformError.With(map[string]string{ + metrics.LabelVertex: isdf.vertexName, + metrics.LabelPipeline: isdf.pipelineName, + metrics.LabelVertexType: string(dfv1.VertexTypeReduceUDF), + metrics.LabelVertexReplicaIndex: strconv.Itoa(int(isdf.vertexReplica)), + }).Inc() + isdf.opts.logger.Errorw("Got an error while invoking WhereTo, dropping the message", zap.Strings("keys", msg.Keys), zap.Error(err)) + continue + } + + if len(to) == 0 { + continue + } + + for _, step := range to { + if _, ok := messagesToStep[step.ToVertexName]; !ok { + messagesToStep[step.ToVertexName] = make([][]isb.Message, len(isdf.toBuffers[step.ToVertexName])) + } + messagesToStep[step.ToVertexName][step.ToVertexPartitionIdx] = append(messagesToStep[step.ToVertexName][step.ToVertexPartitionIdx], msg.Message) + } + + } + return messagesToStep +} + +// writeToBuffer writes to the ISBs. +func (isdf *InterStepDataForward) writeToBufferNew(ctx context.Context, edgeName string, partition int32, resultMessages []isb.Message) ([]isb.Offset, error) { + var ( + writeCount int + writeBytes float64 + ) + + var ISBWriteBackoff = wait.Backoff{ + Steps: math.MaxInt, + Duration: 100 * time.Millisecond, + Factor: 1, + Jitter: 0.1, + } + + writeMessages := resultMessages + + // write to isb with infinite exponential backoff (until shutdown is triggered) + var offsets []isb.Offset + ctxClosedErr := wait.ExponentialBackoff(ISBWriteBackoff, func() (done bool, err error) { + var writeErrs []error + var failedMessages []isb.Message + offsets, writeErrs = isdf.toBuffers[edgeName][partition].Write(ctx, writeMessages) + for i, message := range writeMessages { + writeErr := writeErrs[i] + if writeErr != nil { + // 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(writeErr, &isb.NonRetryableBufferWriteErr{}) { + metrics.DropMessagesCount.With(map[string]string{ + metrics.LabelVertex: isdf.vertexName, + metrics.LabelPipeline: isdf.pipelineName, + metrics.LabelVertexType: string(dfv1.VertexTypeReduceUDF), + metrics.LabelVertexReplicaIndex: strconv.Itoa(int(isdf.vertexReplica)), + metrics.LabelPartitionName: isdf.toBuffers[edgeName][partition].GetName(), + metrics.LabelReason: writeErr.Error(), + }).Inc() + + metrics.DropBytesCount.With(map[string]string{ + metrics.LabelVertex: isdf.vertexName, + metrics.LabelPipeline: isdf.pipelineName, + metrics.LabelVertexType: string(dfv1.VertexTypeReduceUDF), + metrics.LabelVertexReplicaIndex: strconv.Itoa(int(isdf.vertexReplica)), + metrics.LabelPartitionName: isdf.toBuffers[edgeName][partition].GetName(), + metrics.LabelReason: writeErr.Error(), + }).Add(float64(len(message.Payload))) + + isdf.opts.logger.Infow("Dropped message", zap.String("reason", writeErr.Error()), zap.String("vertex", isdf.vertexName), zap.String("pipeline", isdf.pipelineName)) + } else { + failedMessages = append(failedMessages, message) + } + } else { + writeCount++ + writeBytes += float64(len(message.Payload)) + } + } + // retry only the failed messages + if len(failedMessages) > 0 { + isdf.opts.logger.Warnw("Failed to write messages to isb inside pnf", zap.Errors("errors", writeErrs)) + writeMessages = failedMessages + metrics.WriteMessagesError.With(map[string]string{ + metrics.LabelVertex: isdf.vertexName, + metrics.LabelPipeline: isdf.pipelineName, + metrics.LabelVertexType: string(dfv1.VertexTypeReduceUDF), + metrics.LabelVertexReplicaIndex: strconv.Itoa(int(isdf.vertexReplica)), + metrics.LabelPartitionName: isdf.toBuffers[edgeName][partition].GetName()}).Add(float64(len(failedMessages))) + + if ctx.Err() != nil { + // no need to retry if the context is closed + return false, ctx.Err() + } + // keep retrying... + return false, nil + } + return true, nil + }) + + if ctxClosedErr != nil { + isdf.opts.logger.Errorw("Ctx closed while writing messages to ISB", zap.Error(ctxClosedErr)) + return nil, ctxClosedErr + } + metrics.WriteMessagesCount.With(map[string]string{ + metrics.LabelVertex: isdf.vertexName, + metrics.LabelPipeline: isdf.pipelineName, + metrics.LabelVertexType: string(dfv1.VertexTypeReduceUDF), + metrics.LabelVertexReplicaIndex: strconv.Itoa(int(isdf.vertexReplica)), + metrics.LabelPartitionName: isdf.toBuffers[edgeName][partition].GetName()}).Add(float64(writeCount)) + + metrics.WriteBytesCount.With(map[string]string{ + metrics.LabelVertex: isdf.vertexName, + metrics.LabelPipeline: isdf.pipelineName, + metrics.LabelVertexType: string(dfv1.VertexTypeReduceUDF), + metrics.LabelVertexReplicaIndex: strconv.Itoa(int(isdf.vertexReplica)), + metrics.LabelPartitionName: isdf.toBuffers[edgeName][partition].GetName()}).Add(writeBytes) + return offsets, nil } diff --git a/pkg/flatmap/rpc/grpc_flatmap.go b/pkg/flatmap/rpc/grpc_flatmap.go index ac3e21119b..b702847208 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" @@ -66,7 +65,7 @@ func (u *GRPCBasedFlatmap) ApplyMap(ctx context.Context, messageStream <-chan *i // invoke the MapFn method with mapRequestCh channel and send the result to responseCh channel // and any error to errCh channel go func() { - log.Println("MYDEBUG: I'm processing here") + //log.Println("MYDEBUG: I'm processing here") index := 0 resultCh, mapErrCh := u.client.MapFn(ctx, mapRequestCh) for { @@ -81,10 +80,17 @@ func (u *GRPCBasedFlatmap) ApplyMap(ctx context.Context, messageStream <-chan *i } // create a unique message id for each response message which will be used for deduplication index++ - responseCh <- u.parseMapResponse(result) + // TODO(stream): We need to remove the request message from the tracker once this is completed. // As we are streaming messages, we need to have some control field to indicate that this is completed // now, we can do that in the SDK itself. + resp, remove := u.parseMapResponse(result) + if remove { + u.tracker.RemoveRequest(result.Result.GetUuid()) + + } else if resp != nil { + responseCh <- resp + } case err := <-mapErrCh: // TODO(stream): Check error handling here @@ -110,7 +116,7 @@ func (u *GRPCBasedFlatmap) ApplyMap(ctx context.Context, messageStream <-chan *i for { select { case msg, ok := <-messageStream: - log.Println("MYDEBUG: reading for messages here") + //log.Println("MYDEBUG: reading for messages here") // if the requestsStream is closed or if the message is nil, return if !ok || msg == nil { //return @@ -121,7 +127,7 @@ func (u *GRPCBasedFlatmap) ApplyMap(ctx context.Context, messageStream <-chan *i select { // TODO(stream): Check the context end here case mapRequestCh <- d: - log.Println("MYDEBUG: send the message here", d.Uuid) + //log.Println("MYDEBUG: send the message here", d.Uuid) //case <-ctx.Done(): // return } @@ -153,13 +159,18 @@ func (u *GRPCBasedFlatmap) createFlatmapRequest(msg *isb.ReadMessage) *flatmappb return d } -func (u *GRPCBasedFlatmap) parseMapResponse(resp *flatmappb.MapResponse) *types.ResponseFlatmap { +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 { + } + // Request has completed remove from the tracker module + if eor == true { + return nil, true } idx, present := u.tracker.GetIdx(uid) if !present { @@ -182,9 +193,10 @@ func (u *GRPCBasedFlatmap) parseMapResponse(resp *flatmappb.MapResponse) *types. }, Tags: result.GetTags(), } + u.tracker.IncrementRespIdx(uid) return &types.ResponseFlatmap{ ParentMessage: parentRequest, Uid: uid, RespMessage: taggedMessage, - } + }, false } diff --git a/pkg/flatmap/tracker/tracker.go b/pkg/flatmap/tracker/tracker.go index d8dbed2785..c1c28c9ceb 100644 --- a/pkg/flatmap/tracker/tracker.go +++ b/pkg/flatmap/tracker/tracker.go @@ -57,5 +57,9 @@ func (t *Tracker) GetIdx(id string) (int, bool) { return -1, ok } return idx.(int), ok +} +func (t *Tracker) RemoveRequest(id string) { + t.requestMap.Delete(id) + t.responseIdx.Delete(id) } From feec7fb683d73e413336b4242e24e88cfc3f69bc Mon Sep 17 00:00:00 2001 From: Sidhant Kohli Date: Thu, 23 May 2024 12:19:22 -0700 Subject: [PATCH 3/5] init Signed-off-by: Sidhant Kohli --- pkg/flatmap/flatmap_udf.go | 2 +- pkg/flatmap/forward/forward.go | 904 ++++++++++++++++++++------------ pkg/flatmap/rpc/grpc_flatmap.go | 35 +- pkg/flatmap/tracker/tracker.go | 2 +- pkg/isb/message.go | 1 + 5 files changed, 596 insertions(+), 348 deletions(-) diff --git a/pkg/flatmap/flatmap_udf.go b/pkg/flatmap/flatmap_udf.go index 6046b78961..a34935604e 100644 --- a/pkg/flatmap/flatmap_udf.go +++ b/pkg/flatmap/flatmap_udf.go @@ -144,7 +144,7 @@ func (u *FlatMapUDFProcessor) Start(ctx context.Context) error { if err != nil { return fmt.Errorf("failed to create map stream client, %w", err) } - flatmapHandler = rpc.NewUDSgRPCBasedFlatmap(flatMapClient) + flatmapHandler = rpc.NewUDSgRPCBasedFlatmap(flatMapClient, int(*u.VertexInstance.Vertex.Spec.Limits.ReadBatchSize)) // Readiness check if err := flatmapHandler.WaitUntilReady(ctx); err != nil { diff --git a/pkg/flatmap/forward/forward.go b/pkg/flatmap/forward/forward.go index 5f88e0472d..c0b13dcd9a 100644 --- a/pkg/flatmap/forward/forward.go +++ b/pkg/flatmap/forward/forward.go @@ -135,7 +135,7 @@ func NewInterStepDataForward( // Error condition: If there is an error while reading, keep trying // Shutdown flow: On receiving a done signal, we stop the reading process, close the inputMessages channel and return func (isdf *InterStepDataForward) readData() (<-chan *isb.ReadMessage, chan struct{}) { - log := isdf.opts.logger + //log := isdf.opts.logger // inputMessages is the channel on which the data read from the ISB will be written and then // this channel will be consumed by the producer. @@ -168,14 +168,15 @@ func (isdf *InterStepDataForward) readData() (<-chan *isb.ReadMessage, chan stru // 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)))) + //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() } + //inputMessages <- readMessages // Keep streaming the messages read to the inputMessages channel for _, msg := range readMessages { - log.Info("MYDEBUG: streaming in read function ", msg.ReadOffset.String(), " ", time.Now().UnixNano()) + //log.Info("MYDEBUG: streaming in read function ", msg.ReadOffset.String(), " ", time.Now().UnixNano()) inputMessages <- msg } } @@ -184,12 +185,47 @@ func (isdf *InterStepDataForward) readData() (<-chan *isb.ReadMessage, chan stru return inputMessages, stopChan } +func (isdf *InterStepDataForward) processRoutine(inputMessageChan <-chan *types.ResponseFlatmap, writeChan chan<- *types.WriteMsgFlatmap) { + idx := 0 + for msg := range inputMessageChan { + //logger.Info("MYDEBUG: Let's send to resp Chan here") + select { + // TODO(stream): add error handling and shutdown here + default: + idx += 1 + //d := isdf.flatmapUDF.ParseMapResponse(msg) + //taggedMessage := &isb.WriteMessage{ + // Message: isb.Message{ + // Header: isb.Header{ + // MessageInfo: msg.ParentMessage.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-%d", msg.ParentMessage.ReadOffset.String(), idx), + // Keys: msg.RespMessage.Keys, + // }, + // Body: isb.Body{ + // Payload: msg.RespMessage.Payload, + // }, + // }, + // Tags: msg.RespMessage.Tags, + //} + f := &types.WriteMsgFlatmap{ + Message: msg, + AckIt: true, + } + //d := isdf.processWriteMessage(f, true) + //logger.Info("MYDEBUG: Sending to write ", msg.ParentMessage.ReadOffset.String(), " ", time.Now().UnixNano()) + //logger.Info("MYDEBUG: Sending to write", string(msg.RespMessage.Payload), "Sending to write", msg.Uid) + writeChan <- f + } + } +} func (isdf *InterStepDataForward) processUdf(inputMessageChan <-chan *isb.ReadMessage) <-chan *types.WriteMsgFlatmap { ctx := isdf.ctx // 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. - logger := isdf.opts.logger + //logger := isdf.opts.logger // TODO(stream) : enable idle watermark publishing //if len(readMessages) == 0 { @@ -228,25 +264,60 @@ func (isdf *InterStepDataForward) processUdf(inputMessageChan <-chan *isb.ReadMe if err != nil { } - writeChan := make(chan *types.WriteMsgFlatmap) + writeChan := make(chan *types.WriteMsgFlatmap, isdf.opts.readBatchSize) // create a channel which would be passed to the next buffers for writing, // errors in messages/no-acks will be propagated from this as well - go func() { defer close(writeChan) - for msg := range udfRespChan { - //logger.Info("MYDEBUG: Let's send to resp Chan here") - select { - // TODO(stream): add error handling and shutdown here - default: - d := isdf.processWriteMessage(msg, true) - logger.Info("MYDEBUG: Sending to write ", msg.ParentMessage.ReadOffset.String(), " ", time.Now().UnixNano()) - //logger.Info("MYDEBUG: Sending to write", string(msg.RespMessage.Payload), "Sending to write", msg.Uid) - writeChan <- d - } + group := sync.WaitGroup{} + for i := 0; i < 50; i++ { + group.Add(1) + go isdf.processRoutine(udfRespChan, writeChan) } + group.Wait() }() + // + //go func() { + // defer close(writeChan) + // idx := 0 + // for msg := range udfRespChan { + // //logger.Info("MYDEBUG: Let's send to resp Chan here") + // select { + // // TODO(stream): add error handling and shutdown here + // default: + // idx += 1 + // //d := isdf.flatmapUDF.ParseMapResponse() + // taggedMessage := &isb.WriteMessage{ + // Message: isb.Message{ + // Header: isb.Header{ + // MessageInfo: msg.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-%d", msg.ReadOffset.String(), idx), + // Keys: msg.Keys, + // }, + // Body: isb.Body{ + // Payload: msg.Payload, + // }, + // }, + // Tags: msg.Keys, + // } + // f := &types.WriteMsgFlatmap{ + // Message: &types.ResponseFlatmap{ + // ParentMessage: msg, + // Uid: "", + // RespMessage: taggedMessage, + // }, + // AckIt: false, + // } + // //d := isdf.processWriteMessage(f, true) + // //logger.Info("MYDEBUG: Sending to write ", msg.ParentMessage.ReadOffset.String(), " ", time.Now().UnixNano()) + // //logger.Info("MYDEBUG: Sending to write", string(msg.RespMessage.Payload), "Sending to write", msg.Uid) + // writeChan <- f + // } + // } + //}() return writeChan } @@ -313,287 +384,433 @@ type readWriteMessagePair struct { udfError error } -func (isdf *InterStepDataForward) writeRoutine(ctx context.Context, writeMessageCh <-chan *types.WriteMsgFlatmap, ackChan chan<- *types.AckMsgFlatmap) { - flushTimer := time.NewTicker(isdf.opts.flushDuration) - writeMessages := make([]*isb.WriteMessage, 0, isdf.opts.batchSize) - readOffsets := make([]*isb.ReadMessage, 0, isdf.opts.batchSize) - logger := isdf.opts.logger - - // should we flush? - var flush bool - //logger.Info("MYDEBUG: I'm writing to buffer here") - - // error != nil only when the context is closed, so we can safely return (our write loop will try indefinitely - // unless ctx.Done() happens) - flush = true -forwardLoop: - for { - select { - case response, ok := <-writeMessageCh: - if !ok { - // TODO(stream): check the error logic here - break forwardLoop - } - if !response.AckIt { - // TODO(stream): if no ack directly set let's short circuit? - } - - // append the write message to the array - writeMessages = append(writeMessages, response.Message.RespMessage) - readOffsets = append(readOffsets, response.Message.ParentMessage) - - // if the batch size is reached, let's flush - if len(writeMessages) >= isdf.opts.batchSize { - flush = true - } - - case <-flushTimer.C: - // if there are no messages to write, continue - if len(writeMessages) == 0 { - continue - } - - // Since flushTimer is triggered, it is time to flush - flush = true - } - - if flush { - if err := isdf.forwardToBuffers(ctx, &writeMessages); err != nil { - // TODO(stream): mark as no ack directly or retry - } - - for _, offset := range readOffsets { - logger.Info("MYDEBUG: Sending to ack ", offset.ReadOffset.String(), " ", time.Now().UnixNano()) - //logger.Info("MYDEBUG: STEP 1 at ACKING ", offset.ReadOffset) - ackChan <- &types.AckMsgFlatmap{ - Message: offset, - // TODO(stream): should we write for no-ack? - AckIt: true, - } - } - // clear the writeMessages - writeMessages = make([]*isb.WriteMessage, 0, isdf.opts.batchSize) - readOffsets = make([]*isb.ReadMessage, 0, isdf.opts.batchSize) - //flush = false - } - } - - // if there are any messages left, forward them to the ISB - if len(writeMessages) > 0 { - if err := isdf.forwardToBuffers(ctx, &writeMessages); err != nil { - // TODO(stream): mark as no ack directly or retry - } - for _, offset := range readOffsets { - //logger.Info("MYDEBUG: STEP 1 at ACKING ", offset.ReadOffset) - ackChan <- &types.AckMsgFlatmap{ - Message: offset, - // TODO(stream): should we write for no-ack? - AckIt: true, - } - } - // clear the writeMessages - writeMessages = make([]*isb.WriteMessage, 0, isdf.opts.batchSize) - readOffsets = make([]*isb.ReadMessage, 0, isdf.opts.batchSize) - //flush = false - } -} +// +//func (isdf *InterStepDataForward) writeRoutine(ctx context.Context, writeMessageCh <-chan *isb.ReadMessage, ackChan chan<- *isb.ReadMessage, wg *sync.WaitGroup) { +// defer wg.Done() +// //flushTimer := time.NewTicker(3 * isdf.opts.flushDuration) +// writeMessages := make([]*isb.WriteMessage, 0, isdf.opts.readBatchSize) +// readOffsets := make([]*isb.ReadMessage, 0, isdf.opts.readBatchSize) +// logger := isdf.opts.logger +// +// // should we flush? +// var flush bool +// //logger.Info("MYDEBUG: I'm writing to buffer here") +// +// // error != nil only when the context is closed, so we can safely return (our write loop will try indefinitely +// // unless ctx.Done() happens) +// flush = true +// idx := 0 +// logger.Info("MYDEBUG: LATEST WRITE ROUTINE") +//forwardLoop: +// for { +// select { +// case msg, ok := <-writeMessageCh: +// if !ok { +// // TODO(stream): check the error logic here +// break forwardLoop +// } +// //if !response.AckIt { +// // // TODO(stream): if no ack directly set let's short circuit? +// //} +// +// idx += 1 +// taggedMessage := &isb.WriteMessage{ +// Message: isb.Message{ +// Header: isb.Header{ +// MessageInfo: msg.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-%d", msg.ReadOffset.String(), idx), +// Keys: msg.Keys, +// }, +// Body: isb.Body{ +// Payload: msg.Payload, +// }, +// }, +// Tags: msg.Keys, +// } +// //f := &types.WriteMsgFlatmap{ +// // Message: &types.ResponseFlatmap{ +// // ParentMessage: msg, +// // Uid: "", +// // RespMessage: taggedMessage, +// // }, +// // AckIt: true, +// //} +// +// writeMessages = append(writeMessages, taggedMessage) +// readOffsets = append(readOffsets, msg) +// +// // append the write message to the array +// +// // if the batch size is reached, let's flush +// if len(writeMessages) >= int(isdf.opts.readBatchSize) { +// flush = true +// } +// +// //case <-flushTimer.C: +// // // if there are no messages to write, continue +// // if len(writeMessages) == 0 { +// // continue +// // } +// // +// // // Since flushTimer is triggered, it is time to flush +// // flush = true +// } +// +// if flush { +// logger.Info("MYDEBUG: NO WRITE ROUTINE", len(writeMessages)) +// if err := isdf.forwardToBuffers(ctx, &writeMessages); err != nil { +// // TODO(stream): mark as no ack directly or retry +// } +// for _, r := range readOffsets { +// ackChan <- r +// } +// idx = 0 +// //ackChan <- readOffsets +// +// //for _, offset := range readOffsets { +// // logger.Info("MYDEBUG: Sending to ack ", offset.ReadOffset.String(), " ", time.Now().UnixNano()) +// // //logger.Info("MYDEBUG: STEP 1 at ACKING ", offset.ReadOffset) +// // ackChan <- &types.AckMsgFlatmap{ +// // Message: offset, +// // // TODO(stream): should we write for no-ack? +// // AckIt: true, +// // } +// //} +// // clear the writeMessages +// writeMessages = make([]*isb.WriteMessage, 0, isdf.opts.readBatchSize) +// readOffsets = make([]*isb.ReadMessage, 0, isdf.opts.readBatchSize) +// flush = false +// } +// } +// +// // if there are any messages left, forward them to the ISB +// if len(writeMessages) > 0 { +// if err := isdf.forwardToBuffers(ctx, &writeMessages); err != nil { +// // TODO(stream): mark as no ack directly or retry +// } +// for _, r := range readOffsets { +// ackChan <- r +// } +// //for _, offset := range readOffsets { +// // //logger.Info("MYDEBUG: STEP 1 at ACKING ", offset.ReadOffset) +// // ackChan <- &types.AckMsgFlatmap{ +// // Message: offset, +// // // TODO(stream): should we write for no-ack? +// // AckIt: true, +// // } +// //} +// // clear the writeMessages +// writeMessages = make([]*isb.WriteMessage, 0, isdf.opts.readBatchSize) +// readOffsets = make([]*isb.ReadMessage, 0, isdf.opts.readBatchSize) +// flush = false +// } +//} -func (isdf *InterStepDataForward) writeAhead(writeMessageCh <-chan *types.WriteMsgFlatmap) <-chan *types.AckMsgFlatmap { +func (isdf *InterStepDataForward) writeAheadNew(writeMessageCh <-chan *types.WriteMsgFlatmap) <-chan *isb.ReadMessage { // TODO(stream): check buffered channel size //logger := isdf.opts.logger - ackChan := make(chan *types.AckMsgFlatmap) + ackChan := make(chan *isb.ReadMessage, isdf.opts.readBatchSize) ctx := isdf.ctx go func() { defer close(ackChan) group := sync.WaitGroup{} - for i := 0; i < isdf.opts.batchSize; i++ { + for i := 0; i < 4*int(isdf.opts.readBatchSize); i++ { group.Add(1) - go isdf.writeRoutine(ctx, writeMessageCh, ackChan) + go isdf.writeRoutineNew(ctx, writeMessageCh, ackChan, &group) } group.Wait() }() return ackChan +} - //go func() { - // defer close(ackChan) - // flushTimer := time.NewTicker(isdf.opts.flushDuration) - // writeMessages := make([]*isb.WriteMessage, 0, isdf.opts.batchSize) - // readOffsets := make([]*isb.ReadMessage, 0, isdf.opts.batchSize) - // - // // should we flush? - // var flush bool - // //logger.Info("MYDEBUG: I'm writing to buffer here") - // - // // error != nil only when the context is closed, so we can safely return (our write loop will try indefinitely - // // unless ctx.Done() happens) - // flush = true - //forwardLoop: - // for { - // select { - // case response, ok := <-writeMessageCh: - // if !ok { - // // TODO(stream): check the error logic here - // break forwardLoop - // } - // if !response.AckIt { - // // TODO(stream): if no ack directly set let's short circuit? - // } - // - // // append the write message to the array - // writeMessages = append(writeMessages, response.Message.RespMessage) - // readOffsets = append(readOffsets, response.Message.ParentMessage) - // - // // if the batch size is reached, let's flush - // if len(writeMessages) >= isdf.opts.batchSize { - // flush = true - // } - // - // case <-flushTimer.C: - // // if there are no messages to write, continue - // if len(writeMessages) == 0 { - // continue - // } - // - // // Since flushTimer is triggered, it is time to flush - // flush = true - // } - // - // if flush { - // if err := isdf.forwardToBuffers(ctx, &writeMessages); err != nil { - // // TODO(stream): mark as no ack directly or retry - // } - // - // for _, offset := range readOffsets { - // logger.Info("MYDEBUG: Sending to ack ", offset.ReadOffset.String(), " ", time.Now().UnixNano()) - // //logger.Info("MYDEBUG: STEP 1 at ACKING ", offset.ReadOffset) - // ackChan <- &types.AckMsgFlatmap{ - // Message: offset, - // // TODO(stream): should we write for no-ack? - // AckIt: true, - // } - // } - // // clear the writeMessages - // writeMessages = make([]*isb.WriteMessage, 0, isdf.opts.batchSize) - // readOffsets = make([]*isb.ReadMessage, 0, isdf.opts.batchSize) - // //flush = false - // } - // } - // - // // if there are any messages left, forward them to the ISB - // if len(writeMessages) > 0 { - // if err := isdf.forwardToBuffers(ctx, &writeMessages); err != nil { - // // TODO(stream): mark as no ack directly or retry - // } - // for _, offset := range readOffsets { - // //logger.Info("MYDEBUG: STEP 1 at ACKING ", offset.ReadOffset) - // ackChan <- &types.AckMsgFlatmap{ - // Message: offset, - // // TODO(stream): should we write for no-ack? - // AckIt: true, - // } - // } - // // clear the writeMessages - // writeMessages = make([]*isb.WriteMessage, 0, isdf.opts.batchSize) - // readOffsets = make([]*isb.ReadMessage, 0, isdf.opts.batchSize) - // //flush = false - // } - //}() - //go func() { - // logger.Info("MYDEBUG: I'm writing to buffer here") - // defer close(ackChan) - // ctx := isdf.ctx - // // TODO: check pnf.forwardResponses - // // 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) - // var writeOffsets = make(map[string][][]isb.Offset) - // - // for toVertex := range isdf.toBuffers { - // // over allocating to have a predictable pattern - // messageToStep[toVertex] = make([][]isb.Message, len(isdf.toBuffers[toVertex])) - // writeOffsets[toVertex] = make([][]isb.Offset, len(isdf.toBuffers[toVertex])) - // } - // // 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 msg := range writeMessageCh { - // logger.Info("MYDEBUG: Got a msg on writeMessageCh ", msg.Message.Uid) - // if !msg.AckIt { - // // TODO(stream): if no ack directly set let's short circuit? - // } - // parentMessage := msg.Message.ParentMessage - // writeMessage := msg.Message.RespMessage - // - // writeMessage.Headers = parentMessage.Headers - // // add vertex name to the ID, since multiple vertices can publish to the same vertex and we need uniqueness across them - // // TODO(stream): check if this ID is correct for DEDUP? - // logger.Info("MYDEBUG: DEDUP ID ", writeMessage.ID) - // //writeMessage.ID = fmt.Sprintf("%s-%s-%d", parentMessage.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, parentMessage); err != nil { - // // TODO(stream): mark as no ack directly or retry - // //return nil, fmt.Errorf("failed at whereToStep, error: %w", err) - // } - // logger.Info("MYDEBUG: STEP 2 at ACKING ", msg.Message.Uid) - // - // // Forward the message to the edge buffer (could be multiple edges) - // curWriteOffsets, err := isdf.writeToBuffers(ctx, messageToStep) - // if err != nil { - // // TODO(stream): mark as no ack directly or retry - // //return nil, fmt.Errorf("failed to write to toBuffers, error: %w", err) - // } - // logger.Info("MYDEBUG: STEP 3 at ACKING ", msg.Message.Uid) - // // Merge curWriteOffsets into writeOffsets - // for vertexName, toVertexBufferOffsets := range curWriteOffsets { - // for index, offsets := range toVertexBufferOffsets { - // writeOffsets[vertexName][index] = append(writeOffsets[vertexName][index], offsets...) - // } - // } - // logger.Info("MYDEBUG: STEP 4 at ACKING ", len(curWriteOffsets)) - // - // // TODO(stream): Message written successfully -> send for ack - // ackChan <- &types.AckMsgFlatmap{ - // Message: parentMessage, - // AckIt: msg.AckIt, - // } - // logger.Info("MYDEBUG: WRITE sent a msg for ACK ", parentMessage.ReadOffset) - // // TODO(stream): publish new watermark - // } - //}() - //return ackChan +// +//func (isdf *InterStepDataForward) writeAhead(writeMessageCh <-chan *types.WriteMsgFlatmap) <-chan *types.AckMsgFlatmap { +// // TODO(stream): check buffered channel size +// //logger := isdf.opts.logger +// ackChan := make(chan *types.AckMsgFlatmap, isdf.opts.readBatchSize) +// ctx := isdf.ctx +// +// go func() { +// defer close(ackChan) +// group := sync.WaitGroup{} +// for i := 0; i < 1; i++ { +// group.Add(1) +// go isdf.writeRoutine(ctx, writeMessageCh, ackChan) +// } +// group.Wait() +// }() +// return ackChan +// +// //go func() { +// // defer close(ackChan) +// // flushTimer := time.NewTicker(isdf.opts.flushDuration) +// // writeMessages := make([]*isb.WriteMessage, 0, isdf.opts.batchSize) +// // readOffsets := make([]*isb.ReadMessage, 0, isdf.opts.batchSize) +// // +// // // should we flush? +// // var flush bool +// // //logger.Info("MYDEBUG: I'm writing to buffer here") +// // +// // // error != nil only when the context is closed, so we can safely return (our write loop will try indefinitely +// // // unless ctx.Done() happens) +// // flush = true +// //forwardLoop: +// // for { +// // select { +// // case response, ok := <-writeMessageCh: +// // if !ok { +// // // TODO(stream): check the error logic here +// // break forwardLoop +// // } +// // if !response.AckIt { +// // // TODO(stream): if no ack directly set let's short circuit? +// // } +// // +// // // append the write message to the array +// // writeMessages = append(writeMessages, response.Message.RespMessage) +// // readOffsets = append(readOffsets, response.Message.ParentMessage) +// // +// // // if the batch size is reached, let's flush +// // if len(writeMessages) >= isdf.opts.batchSize { +// // flush = true +// // } +// // +// // case <-flushTimer.C: +// // // if there are no messages to write, continue +// // if len(writeMessages) == 0 { +// // continue +// // } +// // +// // // Since flushTimer is triggered, it is time to flush +// // flush = true +// // } +// // +// // if flush { +// // if err := isdf.forwardToBuffers(ctx, &writeMessages); err != nil { +// // // TODO(stream): mark as no ack directly or retry +// // } +// // +// // for _, offset := range readOffsets { +// // logger.Info("MYDEBUG: Sending to ack ", offset.ReadOffset.String(), " ", time.Now().UnixNano()) +// // //logger.Info("MYDEBUG: STEP 1 at ACKING ", offset.ReadOffset) +// // ackChan <- &types.AckMsgFlatmap{ +// // Message: offset, +// // // TODO(stream): should we write for no-ack? +// // AckIt: true, +// // } +// // } +// // // clear the writeMessages +// // writeMessages = make([]*isb.WriteMessage, 0, isdf.opts.batchSize) +// // readOffsets = make([]*isb.ReadMessage, 0, isdf.opts.batchSize) +// // //flush = false +// // } +// // } +// // +// // // if there are any messages left, forward them to the ISB +// // if len(writeMessages) > 0 { +// // if err := isdf.forwardToBuffers(ctx, &writeMessages); err != nil { +// // // TODO(stream): mark as no ack directly or retry +// // } +// // for _, offset := range readOffsets { +// // //logger.Info("MYDEBUG: STEP 1 at ACKING ", offset.ReadOffset) +// // ackChan <- &types.AckMsgFlatmap{ +// // Message: offset, +// // // TODO(stream): should we write for no-ack? +// // AckIt: true, +// // } +// // } +// // // clear the writeMessages +// // writeMessages = make([]*isb.WriteMessage, 0, isdf.opts.batchSize) +// // readOffsets = make([]*isb.ReadMessage, 0, isdf.opts.batchSize) +// // //flush = false +// // } +// //}() +// //go func() { +// // logger.Info("MYDEBUG: I'm writing to buffer here") +// // defer close(ackChan) +// // ctx := isdf.ctx +// // // TODO: check pnf.forwardResponses +// // // 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) +// // var writeOffsets = make(map[string][][]isb.Offset) +// // +// // for toVertex := range isdf.toBuffers { +// // // over allocating to have a predictable pattern +// // messageToStep[toVertex] = make([][]isb.Message, len(isdf.toBuffers[toVertex])) +// // writeOffsets[toVertex] = make([][]isb.Offset, len(isdf.toBuffers[toVertex])) +// // } +// // // 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 msg := range writeMessageCh { +// // logger.Info("MYDEBUG: Got a msg on writeMessageCh ", msg.Message.Uid) +// // if !msg.AckIt { +// // // TODO(stream): if no ack directly set let's short circuit? +// // } +// // parentMessage := msg.Message.ParentMessage +// // writeMessage := msg.Message.RespMessage +// // +// // writeMessage.Headers = parentMessage.Headers +// // // add vertex name to the ID, since multiple vertices can publish to the same vertex and we need uniqueness across them +// // // TODO(stream): check if this ID is correct for DEDUP? +// // logger.Info("MYDEBUG: DEDUP ID ", writeMessage.ID) +// // //writeMessage.ID = fmt.Sprintf("%s-%s-%d", parentMessage.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, parentMessage); err != nil { +// // // TODO(stream): mark as no ack directly or retry +// // //return nil, fmt.Errorf("failed at whereToStep, error: %w", err) +// // } +// // logger.Info("MYDEBUG: STEP 2 at ACKING ", msg.Message.Uid) +// // +// // // Forward the message to the edge buffer (could be multiple edges) +// // curWriteOffsets, err := isdf.writeToBuffers(ctx, messageToStep) +// // if err != nil { +// // // TODO(stream): mark as no ack directly or retry +// // //return nil, fmt.Errorf("failed to write to toBuffers, error: %w", err) +// // } +// // logger.Info("MYDEBUG: STEP 3 at ACKING ", msg.Message.Uid) +// // // Merge curWriteOffsets into writeOffsets +// // for vertexName, toVertexBufferOffsets := range curWriteOffsets { +// // for index, offsets := range toVertexBufferOffsets { +// // writeOffsets[vertexName][index] = append(writeOffsets[vertexName][index], offsets...) +// // } +// // } +// // logger.Info("MYDEBUG: STEP 4 at ACKING ", len(curWriteOffsets)) +// // +// // // TODO(stream): Message written successfully -> send for ack +// // ackChan <- &types.AckMsgFlatmap{ +// // Message: parentMessage, +// // AckIt: msg.AckIt, +// // } +// // logger.Info("MYDEBUG: WRITE sent a msg for ACK ", parentMessage.ReadOffset) +// // // TODO(stream): publish new watermark +// // } +// //}() +// //return ackChan +// +//} -} +//func (isdf *InterStepDataForward) ackRoutine(ctx context.Context, ackMsgChan <-chan *types.AckMsgFlatmap) { +// logger := isdf.opts.logger +//forwardLoop: +// for { +// select { +// case response, ok := <-ackMsgChan: +// if !ok { +// break forwardLoop +// } +// if response.AckIt { +// ackMessages := []isb.Offset{response.Message.ReadOffset} +// if err := isdf.ackFromBuffer(ctx, ackMessages); err != nil { +// // TODO(stream): we have retried in the ackFromBuffer, should we trigger +// // shutdown here then? +// } +// logger.Info("MYDEBUG: Done with ack ", ackMessages, " ", time.Now().UnixNano()) +// } else { +// noAckMessages := []isb.Offset{response.Message.ReadOffset} +// isdf.noAckMessages(ctx, noAckMessages) +// } +// } +// } +//} -func (isdf *InterStepDataForward) ackRoutine(ctx context.Context, ackMsgChan <-chan *types.AckMsgFlatmap) { +func (isdf *InterStepDataForward) ackRoutine(ctx context.Context, ackMsgChan <-chan *isb.ReadMessage) { logger := isdf.opts.logger + logger.Info("MYDEBUG: NO WG ACK ROUTINE ", isdf.opts.readBatchSize) forwardLoop: for { select { + case <-ctx.Done(): + break forwardLoop + case response, ok := <-ackMsgChan: if !ok { + logger.Info("MYDEBUG: ERROR ACK ROUTINE") break forwardLoop } - if response.AckIt { - ackMessages := []isb.Offset{response.Message.ReadOffset} - if err := isdf.ackFromBuffer(ctx, ackMessages); err != nil { - // TODO(stream): we have retried in the ackFromBuffer, should we trigger - // shutdown here then? - } - logger.Info("MYDEBUG: Done with ack ", ackMessages, " ", time.Now().UnixNano()) - } else { - noAckMessages := []isb.Offset{response.Message.ReadOffset} - isdf.noAckMessages(ctx, noAckMessages) + ackMessages := []isb.Offset{response.ReadOffset} + if err := isdf.ackFromBuffer(ctx, ackMessages); err != nil { + // TODO(stream): we have retried in the ackFromBuffer, should we trigger + // shutdown here then? } } - } +} + +func (isdf *InterStepDataForward) writeRoutineNew(ctx context.Context, writeMessageCh <-chan *types.WriteMsgFlatmap, ackChan chan<- *isb.ReadMessage, wg *sync.WaitGroup) { + logger := isdf.opts.logger + logger.Info("PLEASE WORK WRITE ROUTINE") + ////flushTimer := time.NewTicker(isdf.opts.flushDuration) + //writeMessages := make([]*isb.WriteMessage, 0, isdf.opts.batchSize) + //readOffsets := make([]*isb.ReadMessage, 0, isdf.opts.batchSize) + defer wg.Done() +forwardLoop: + for { + idx := 0 + select { + case msg, ok := <-writeMessageCh: + if !ok { + logger.Info("MYDEBUG: ERROR WRITE ROUTINE") + // TODO(stream): check the error logic here + break forwardLoop + } + idx += 1 + //taggedMessage := &isb.WriteMessage{ + // Message: isb.Message{ + // Header: isb.Header{ + // MessageInfo: msg.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-%d", msg.ReadOffset.String(), idx), + // Keys: msg.Keys, + // }, + // Body: isb.Body{ + // Payload: msg.Payload, + // }, + // ReadOffset: msg.ReadOffset, + // }, + // Tags: msg.Keys, + //} + + taggedMessage := msg.Message.RespMessage + 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])) + } + //writeMessages = append(writeMessages, taggedMessage) + //readOffsets = append(readOffsets, msg) + // + //if len(writeMessages) >= int(isdf.opts.readBatchSize) { + // logger.Info("MYDEBUG: PLEASE WORK ROUTINE", len(writeMessages)) + writeMessages := []*isb.WriteMessage{taggedMessage} + if err := isdf.forwardToBuffersNew(ctx, &writeMessages, msg.Message.ParentMessage, messageToStep); err != nil { + // // TODO(stream): mark as no ack directly or retry + } + // for _, offset := range readOffsets { + // ackChan <- offset + // } + // writeMessages = make([]*isb.WriteMessage, 0, isdf.opts.batchSize) + // readOffsets = make([]*isb.ReadMessage, 0, isdf.opts.batchSize) + //} + //writeMessages := []*isb.WriteMessage{taggedMessage} + //if err := isdf.forwardToBuffers(ctx, &writeMessages); err != nil { + // // // TODO(stream): mark as no ack directly or retry + //} + ackChan <- msg.Message.ParentMessage + } + } } +// //func (isdf *InterStepDataForward) ackPrevBuffer(ackMsgChan <-chan *types.AckMsgFlatmap) { // ctx := isdf.ctx // flushTimer := time.NewTicker(isdf.opts.flushDuration) @@ -681,7 +898,7 @@ forwardLoop: // isdf.noAckMessages(ctx, noAckMessages) // noAckMessages = make([]isb.Offset, 0, isdf.opts.batchSize) // } - +// //} // forwardAChunk forwards a chunk of message from the fromBufferPartition to the toBuffers. It does the Read -> Process -> Forward -> Ack chain @@ -698,17 +915,16 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { // processing data writeMsgChan := isdf.processUdf(inputMessagesChan) - // writing data to next ISB buffer - ackMsgChan := isdf.writeAhead(writeMsgChan) + //writing data to next ISB buffer + ackMsgChan := isdf.writeAheadNew(writeMsgChan) + //ackMsgChan := isdf.writeAheadNew(inputMessagesChan) // Ack to previous ISB go func() { - group := sync.WaitGroup{} - for i := 0; i < isdf.opts.batchSize; i++ { - group.Add(1) + for i := 0; i < 4*int(isdf.opts.readBatchSize); i++ { + //go isdf.ackRoutine(isdf.ctx, ackMsgChan) go isdf.ackRoutine(isdf.ctx, ackMsgChan) } - group.Wait() }() // TODO(stream): check ideal way to wait here @@ -716,7 +932,7 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { } // 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 { +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 { @@ -734,7 +950,7 @@ func (isdf *InterStepDataForward) whereToStep(writeMessage *isb.WriteMessage, me 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] = append(messageToStep[t.ToVertexName][t.ToVertexPartitionIdx], writeMessage.Message) + messageToStep[t.ToVertexName][t.ToVertexPartitionIdx] = writeMessage.Message } return nil } @@ -742,7 +958,7 @@ func (isdf *InterStepDataForward) whereToStep(writeMessage *isb.WriteMessage, me // 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, + 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. @@ -762,66 +978,71 @@ func (isdf *InterStepDataForward) writeToBuffers( } // 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, messages []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 + //totalCount int writeCount int writeBytes float64 ) - totalCount = len(messages) - writeOffsets = make([]isb.Offset, 0, totalCount) + //totalCount = len(messages) + //writeOffsets = make([]isb.Offset, 0, totalCount) for { - _writeOffsets, errs := toBufferPartition.Write(ctx, messages) + // 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 + var failedMessages isb.Message needRetry := false - for idx, msg := range messages { - if err = errs[idx]; 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 = append(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:%d, %v", len(failedMessages), errs) - } - } + //for idx, msg := range messages { + // EXTRA + if err != nil { + //if err = errs[idx]; 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 { - writeCount++ - writeBytes += float64(len(msg.Payload)) - // we support write offsets only for jetstream - if _writeOffsets != nil { - writeOffsets = append(writeOffsets, _writeOffsets[idx]) + 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", @@ -831,7 +1052,7 @@ func (isdf *InterStepDataForward) writeToBuffer(ctx context.Context, toBufferPar zap.String(metrics.LabelPartitionName, toBufferPartition.GetName()), ) // set messages to failed for the retry - messages = failedMessages + msg = failedMessages // TODO: implement retry with backoff etc. time.Sleep(isdf.opts.retryInterval) } else { @@ -916,6 +1137,26 @@ func (isdf *InterStepDataForward) processWriteMessage(msg *types.ResponseFlatmap } } +func (isdf *InterStepDataForward) forwardToBuffersNew(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 +} + // forwardToBuffers writes the messages to the ISBs concurrently for each partition. func (isdf *InterStepDataForward) forwardToBuffers(ctx context.Context, writeMessages *[]*isb.WriteMessage) error { if len(*writeMessages) == 0 { @@ -1008,13 +1249,16 @@ func (isdf *InterStepDataForward) writeToBufferNew(ctx context.Context, edgeName writeMessages := resultMessages // write to isb with infinite exponential backoff (until shutdown is triggered) - var offsets []isb.Offset + //var offsets []isb.Offset + offsets := make([]isb.Offset, len(writeMessages)) ctxClosedErr := wait.ExponentialBackoff(ISBWriteBackoff, func() (done bool, err error) { var writeErrs []error var failedMessages []isb.Message - offsets, writeErrs = isdf.toBuffers[edgeName][partition].Write(ctx, writeMessages) - for i, message := range writeMessages { - writeErr := writeErrs[i] + //offsets, writeErrs = isdf.toBuffers[edgeName][partition].Write(ctx, writeMessages) + for _, message := range writeMessages { + offsets = append(offsets, message.ReadOffset) + var writeErr error = nil + //writeErr := writeErrs[i] if writeErr != nil { // 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 diff --git a/pkg/flatmap/rpc/grpc_flatmap.go b/pkg/flatmap/rpc/grpc_flatmap.go index b702847208..458fbb9971 100644 --- a/pkg/flatmap/rpc/grpc_flatmap.go +++ b/pkg/flatmap/rpc/grpc_flatmap.go @@ -19,12 +19,14 @@ 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 + client flatmapper.Client + tracker *tracker.Tracker + readBatchSize int + idx int } -func NewUDSgRPCBasedFlatmap(client flatmapper.Client) *GRPCBasedFlatmap { - return &GRPCBasedFlatmap{client: client, tracker: tracker.NewTracker()} +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. @@ -58,8 +60,8 @@ func (u *GRPCBasedFlatmap) WaitUntilReady(ctx context.Context) error { func (u *GRPCBasedFlatmap) ApplyMap(ctx context.Context, messageStream <-chan *isb.ReadMessage) (<-chan *types.ResponseFlatmap, <-chan error) { var ( errCh = make(chan error) - responseCh = make(chan *types.ResponseFlatmap) - mapRequestCh = make(chan *flatmappb.MapRequest) + responseCh = make(chan *types.ResponseFlatmap, u.readBatchSize) + mapRequestCh = make(chan *flatmappb.MapRequest, u.readBatchSize) ) // invoke the MapFn method with mapRequestCh channel and send the result to responseCh channel @@ -84,9 +86,9 @@ func (u *GRPCBasedFlatmap) ApplyMap(ctx context.Context, messageStream <-chan *i // TODO(stream): We need to remove the request message from the tracker once this is completed. // As we are streaming messages, we need to have some control field to indicate that this is completed // now, we can do that in the SDK itself. - resp, remove := u.parseMapResponse(result) + resp, remove := u.ParseMapResponse(result) if remove { - u.tracker.RemoveRequest(result.Result.GetUuid()) + //u.tracker.RemoveRequest(result.Result.GetUuid()) } else if resp != nil { responseCh <- resp @@ -159,7 +161,7 @@ func (u *GRPCBasedFlatmap) createFlatmapRequest(msg *isb.ReadMessage) *flatmappb return d } -func (u *GRPCBasedFlatmap) parseMapResponse(resp *flatmappb.MapResponse) (parsedResp *types.ResponseFlatmap, requestDone bool) { +func (u *GRPCBasedFlatmap) ParseMapResponse(resp *flatmappb.MapResponse) (parsedResp *types.ResponseFlatmap, requestDone bool) { result := resp.Result eor := result.GetEOR() uid := result.GetUuid() @@ -172,11 +174,11 @@ func (u *GRPCBasedFlatmap) parseMapResponse(resp *flatmappb.MapResponse) (parsed if eor == true { return nil, true } - idx, present := u.tracker.GetIdx(uid) - if !present { - u.tracker.NewResponse(uid) - idx = 1 - } + //idx, present := u.tracker.GetIdx(uid) + //if !present { + // u.tracker.NewResponse(uid) + // idx = 1 + //} keys := result.GetKeys() taggedMessage := &isb.WriteMessage{ Message: isb.Message{ @@ -184,7 +186,7 @@ func (u *GRPCBasedFlatmap) parseMapResponse(resp *flatmappb.MapResponse) (parsed 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-%d", parentRequest.ReadOffset.String(), idx), + ID: fmt.Sprintf("%s-%d", parentRequest.ReadOffset.String(), u.idx), Keys: keys, }, Body: isb.Body{ @@ -193,7 +195,8 @@ func (u *GRPCBasedFlatmap) parseMapResponse(resp *flatmappb.MapResponse) (parsed }, Tags: result.GetTags(), } - u.tracker.IncrementRespIdx(uid) + u.idx += 1 + //u.tracker.IncrementRespIdx(uid) return &types.ResponseFlatmap{ ParentMessage: parentRequest, Uid: uid, diff --git a/pkg/flatmap/tracker/tracker.go b/pkg/flatmap/tracker/tracker.go index c1c28c9ceb..5443b967bc 100644 --- a/pkg/flatmap/tracker/tracker.go +++ b/pkg/flatmap/tracker/tracker.go @@ -28,7 +28,7 @@ func GetNewId() string { func (t *Tracker) AddRequest(msg *isb.ReadMessage) string { id := GetNewId() - t.requestMap.Store(id, msg) + //t.requestMap.Store(id, msg) return id } diff --git a/pkg/isb/message.go b/pkg/isb/message.go index fc475ccd7e..a388fad424 100644 --- a/pkg/isb/message.go +++ b/pkg/isb/message.go @@ -82,6 +82,7 @@ type Body struct { type Message struct { Header Body + ReadOffset Offset } // ReadMessage is the message read from the buffer. From e50197adf27c3c0c9ba470e5e6061d6415b08a8d Mon Sep 17 00:00:00 2001 From: Sidhant Kohli Date: Thu, 23 May 2024 12:34:39 -0700 Subject: [PATCH 4/5] init Signed-off-by: Sidhant Kohli --- pkg/flatmap/forward/forward.go | 24 ++++++++++++++---------- pkg/flatmap/rpc/grpc_flatmap.go | 2 +- pkg/flatmap/tracker/tracker.go | 4 ++-- 3 files changed, 17 insertions(+), 13 deletions(-) diff --git a/pkg/flatmap/forward/forward.go b/pkg/flatmap/forward/forward.go index c0b13dcd9a..b5abf3c91e 100644 --- a/pkg/flatmap/forward/forward.go +++ b/pkg/flatmap/forward/forward.go @@ -185,11 +185,14 @@ func (isdf *InterStepDataForward) readData() (<-chan *isb.ReadMessage, chan stru return inputMessages, stopChan } -func (isdf *InterStepDataForward) processRoutine(inputMessageChan <-chan *types.ResponseFlatmap, writeChan chan<- *types.WriteMsgFlatmap) { +func (isdf *InterStepDataForward) processRoutine(ctx context.Context, inputMessageChan <-chan *types.ResponseFlatmap, writeChan chan<- *types.WriteMsgFlatmap) { + defer close(writeChan) idx := 0 for msg := range inputMessageChan { //logger.Info("MYDEBUG: Let's send to resp Chan here") select { + case <-ctx.Done(): + return // TODO(stream): add error handling and shutdown here default: idx += 1 @@ -268,15 +271,16 @@ func (isdf *InterStepDataForward) processUdf(inputMessageChan <-chan *isb.ReadMe // create a channel which would be passed to the next buffers for writing, // errors in messages/no-acks will be propagated from this as well - go func() { - defer close(writeChan) - group := sync.WaitGroup{} - for i := 0; i < 50; i++ { - group.Add(1) - go isdf.processRoutine(udfRespChan, writeChan) - } - group.Wait() - }() + go isdf.processRoutine(ctx, udfRespChan, writeChan) + //go func() { + // defer close(writeChan) + // group := sync.WaitGroup{} + // for i := 0; i < 50; i++ { + // group.Add(1) + // go isdf.processRoutine(udfRespChan, writeChan) + // } + // group.Wait() + //}() // //go func() { // defer close(writeChan) diff --git a/pkg/flatmap/rpc/grpc_flatmap.go b/pkg/flatmap/rpc/grpc_flatmap.go index 458fbb9971..71e78589f0 100644 --- a/pkg/flatmap/rpc/grpc_flatmap.go +++ b/pkg/flatmap/rpc/grpc_flatmap.go @@ -88,7 +88,7 @@ func (u *GRPCBasedFlatmap) ApplyMap(ctx context.Context, messageStream <-chan *i // now, we can do that in the SDK itself. resp, remove := u.ParseMapResponse(result) if remove { - //u.tracker.RemoveRequest(result.Result.GetUuid()) + u.tracker.RemoveRequest(result.Result.GetUuid()) } else if resp != nil { responseCh <- resp diff --git a/pkg/flatmap/tracker/tracker.go b/pkg/flatmap/tracker/tracker.go index 5443b967bc..24d8b76499 100644 --- a/pkg/flatmap/tracker/tracker.go +++ b/pkg/flatmap/tracker/tracker.go @@ -28,7 +28,7 @@ func GetNewId() string { func (t *Tracker) AddRequest(msg *isb.ReadMessage) string { id := GetNewId() - //t.requestMap.Store(id, msg) + t.requestMap.Store(id, msg) return id } @@ -61,5 +61,5 @@ func (t *Tracker) GetIdx(id string) (int, bool) { func (t *Tracker) RemoveRequest(id string) { t.requestMap.Delete(id) - t.responseIdx.Delete(id) + //t.responseIdx.Delete(id) } From ca2c3ef31a9c56a87df0ce7fa28b6e0fc493b81a Mon Sep 17 00:00:00 2001 From: Sidhant Kohli Date: Thu, 30 May 2024 15:26:09 -0700 Subject: [PATCH 5/5] init Signed-off-by: Sidhant Kohli --- pkg/flatmap/forward/forward.go | 203 +++----------------------------- pkg/flatmap/rpc/grpc_flatmap.go | 138 ++++++++++++++-------- pkg/flatmap/tracker/tracker.go | 31 ++++- 3 files changed, 133 insertions(+), 239 deletions(-) diff --git a/pkg/flatmap/forward/forward.go b/pkg/flatmap/forward/forward.go index b5abf3c91e..8df802e459 100644 --- a/pkg/flatmap/forward/forward.go +++ b/pkg/flatmap/forward/forward.go @@ -173,11 +173,16 @@ func (isdf *InterStepDataForward) readData() (<-chan *isb.ReadMessage, chan stru 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() } + //inputMessages <- readMessages // Keep streaming the messages read to the inputMessages channel for _, msg := range readMessages { + if msg.Kind == isb.Data { + inputMessages <- msg + 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()}).Inc() + //dataMessages = append(dataMessages, m) + } //log.Info("MYDEBUG: streaming in read function ", msg.ReadOffset.String(), " ", time.Now().UnixNano()) - inputMessages <- msg } } } @@ -188,8 +193,10 @@ func (isdf *InterStepDataForward) readData() (<-chan *isb.ReadMessage, chan stru func (isdf *InterStepDataForward) processRoutine(ctx context.Context, inputMessageChan <-chan *types.ResponseFlatmap, writeChan chan<- *types.WriteMsgFlatmap) { defer close(writeChan) idx := 0 + //logger := isdf.opts.logger + //timer := time.Now() for msg := range inputMessageChan { - //logger.Info("MYDEBUG: Let's send to resp Chan here") + //logger.Info("MYDEBUG: Let's send to resp Chan here ", idx, "::", time.Now()) select { case <-ctx.Done(): return @@ -230,38 +237,7 @@ func (isdf *InterStepDataForward) processUdf(inputMessageChan <-chan *isb.ReadMe //logger := isdf.opts.logger // TODO(stream) : enable idle watermark publishing - - //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 - //} - //logger.Info("MYDEBUG: I'm processing here") - - // Send thr requests to the grpc server for results, which are received in the udfRespChan + // Send the requests to the grpc server for results, which are received in the udfRespChan udfRespChan, err := isdf.flatmapUDF.ApplyMap(ctx, inputMessageChan) // TODO(stream): check error handling if err != nil { @@ -272,56 +248,6 @@ func (isdf *InterStepDataForward) processUdf(inputMessageChan <-chan *isb.ReadMe // create a channel which would be passed to the next buffers for writing, // errors in messages/no-acks will be propagated from this as well go isdf.processRoutine(ctx, udfRespChan, writeChan) - //go func() { - // defer close(writeChan) - // group := sync.WaitGroup{} - // for i := 0; i < 50; i++ { - // group.Add(1) - // go isdf.processRoutine(udfRespChan, writeChan) - // } - // group.Wait() - //}() - // - //go func() { - // defer close(writeChan) - // idx := 0 - // for msg := range udfRespChan { - // //logger.Info("MYDEBUG: Let's send to resp Chan here") - // select { - // // TODO(stream): add error handling and shutdown here - // default: - // idx += 1 - // //d := isdf.flatmapUDF.ParseMapResponse() - // taggedMessage := &isb.WriteMessage{ - // Message: isb.Message{ - // Header: isb.Header{ - // MessageInfo: msg.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-%d", msg.ReadOffset.String(), idx), - // Keys: msg.Keys, - // }, - // Body: isb.Body{ - // Payload: msg.Payload, - // }, - // }, - // Tags: msg.Keys, - // } - // f := &types.WriteMsgFlatmap{ - // Message: &types.ResponseFlatmap{ - // ParentMessage: msg, - // Uid: "", - // RespMessage: taggedMessage, - // }, - // AckIt: false, - // } - // //d := isdf.processWriteMessage(f, true) - // //logger.Info("MYDEBUG: Sending to write ", msg.ParentMessage.ReadOffset.String(), " ", time.Now().UnixNano()) - // //logger.Info("MYDEBUG: Sending to write", string(msg.RespMessage.Payload), "Sending to write", msg.Uid) - // writeChan <- f - // } - // } - //}() return writeChan } @@ -521,7 +447,7 @@ func (isdf *InterStepDataForward) writeAheadNew(writeMessageCh <-chan *types.Wri go func() { defer close(ackChan) group := sync.WaitGroup{} - for i := 0; i < 4*int(isdf.opts.readBatchSize); i++ { + for i := 0; i < int(isdf.opts.readBatchSize); i++ { group.Add(1) go isdf.writeRoutineNew(ctx, writeMessageCh, ackChan, &group) } @@ -814,97 +740,6 @@ forwardLoop: } } -// -//func (isdf *InterStepDataForward) ackPrevBuffer(ackMsgChan <-chan *types.AckMsgFlatmap) { -// ctx := isdf.ctx -// flushTimer := time.NewTicker(isdf.opts.flushDuration) -// ackMessages := make([]isb.Offset, 0, isdf.opts.batchSize) -// noAckMessages := make([]isb.Offset, 0, isdf.opts.batchSize) -// logger := isdf.opts.logger -// -// // should we flush? -// var flushAck bool -// -// // should we flush? -// var flushNoAck bool -// -// flushAck = true -// flushNoAck = true -// // error != nil only when the context is closed, so we can safely return (our write loop will try indefinitely -// // unless ctx.Done() happens) -//forwardLoop: -// for { -// select { -// case response, ok := <-ackMsgChan: -// if !ok { -// break forwardLoop -// } -// -// if response.AckIt { -// // append the ack message to the array -// ackMessages = append(ackMessages, response.Message.ReadOffset) -// } else { -// // append the ack message to the array -// noAckMessages = append(noAckMessages, response.Message.ReadOffset) -// -// } -// -// // if the batch size is reached, let's flush -// if len(ackMessages) >= isdf.opts.batchSize { -// flushAck = true -// } -// -// // if the batch size is reached, let's flush -// if len(noAckMessages) >= isdf.opts.batchSize { -// flushNoAck = true -// } -// -// case <-flushTimer.C: -// // if there are no messages to write, continue -// if len(ackMessages) == 0 { -// continue -// } -// -// // Since flushTimer is triggered, it is time to flush -// flushAck = true -// flushNoAck = true -// } -// -// if flushAck { -// if err := isdf.ackFromBuffer(ctx, ackMessages); err != nil { -// // TODO(stream): we have retried in the ackFromBuffer, should we trigger -// // shutdown here then? -// } -// logger.Info("MYDEBUG: Sending to ack ", ackMessages, " ", time.Now().UnixNano()) -// ackMessages = make([]isb.Offset, 0, isdf.opts.batchSize) -// //flushAck = false -// } -// -// if flushNoAck { -// isdf.noAckMessages(ctx, noAckMessages) -// noAckMessages = make([]isb.Offset, 0, isdf.opts.batchSize) -// //flushNoAck = false -// } -// } -// -// // if there are any messages left, forward them to the ISB -// if len(ackMessages) > 0 { -// if err := isdf.ackFromBuffer(ctx, ackMessages); err != nil { -// // TODO(stream): we have retried in the ackFromBuffer, should we trigger -// // shutdown here then? -// //return -// } -// ackMessages = make([]isb.Offset, 0, isdf.opts.batchSize) -// } -// -// // if there are any messages left, forward them to the ISB -// if len(noAckMessages) > 0 { -// isdf.noAckMessages(ctx, noAckMessages) -// noAckMessages = make([]isb.Offset, 0, isdf.opts.batchSize) -// } -// -//} - // 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, @@ -925,7 +760,7 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) { // Ack to previous ISB go func() { - for i := 0; i < 4*int(isdf.opts.readBatchSize); i++ { + for i := 0; i < int(isdf.opts.readBatchSize); i++ { //go isdf.ackRoutine(isdf.ctx, ackMsgChan) go isdf.ackRoutine(isdf.ctx, ackMsgChan) } @@ -993,9 +828,9 @@ func (isdf *InterStepDataForward) writeToBuffer(ctx context.Context, toBufferPar for { // EXTRA - var _writeOffsets []isb.Offset = nil - var errs []error = nil - //_writeOffsets, errs := toBufferPartition.Write(ctx, []isb.Message{msg}) + //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 @@ -1258,11 +1093,11 @@ func (isdf *InterStepDataForward) writeToBufferNew(ctx context.Context, edgeName ctxClosedErr := wait.ExponentialBackoff(ISBWriteBackoff, func() (done bool, err error) { var writeErrs []error var failedMessages []isb.Message - //offsets, writeErrs = isdf.toBuffers[edgeName][partition].Write(ctx, writeMessages) - for _, message := range writeMessages { + offsets, writeErrs = isdf.toBuffers[edgeName][partition].Write(ctx, writeMessages) + for i, message := range writeMessages { offsets = append(offsets, message.ReadOffset) - var writeErr error = nil - //writeErr := writeErrs[i] + //var writeErr error = nil + writeErr := writeErrs[i] if writeErr != nil { // 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 diff --git a/pkg/flatmap/rpc/grpc_flatmap.go b/pkg/flatmap/rpc/grpc_flatmap.go index 71e78589f0..0319f48212 100644 --- a/pkg/flatmap/rpc/grpc_flatmap.go +++ b/pkg/flatmap/rpc/grpc_flatmap.go @@ -2,6 +2,7 @@ package rpc import ( "fmt" + "log" "time" flatmappb "github.com/numaproj/numaflow-go/pkg/apis/proto/flatmap/v1" @@ -41,7 +42,7 @@ func (u *GRPCBasedFlatmap) CloseConn(ctx context.Context) error { // WaitUntilReady waits until the reduce udf is connected. func (u *GRPCBasedFlatmap) WaitUntilReady(ctx context.Context) error { - log := logging.FromContext(ctx) + logger := logging.FromContext(ctx) for { select { case <-ctx.Done(): @@ -50,7 +51,7 @@ func (u *GRPCBasedFlatmap) WaitUntilReady(ctx context.Context) error { if _, err := u.client.IsReady(ctx, &emptypb.Empty{}); err == nil { return nil } else { - log.Infof("waiting for reduce udf to be ready: %v", err) + logger.Infof("waiting for reduce udf to be ready: %v", err) time.Sleep(1 * time.Second) } } @@ -64,50 +65,51 @@ func (u *GRPCBasedFlatmap) ApplyMap(ctx context.Context, messageStream <-chan *i mapRequestCh = make(chan *flatmappb.MapRequest, u.readBatchSize) ) + log.Println("MYDEBUG: NO PROCESS WITH UDF") // invoke the MapFn method with mapRequestCh channel and send the result to responseCh channel // and any error to errCh channel - go func() { - //log.Println("MYDEBUG: I'm processing here") - index := 0 - resultCh, mapErrCh := u.client.MapFn(ctx, mapRequestCh) - for { - select { - case result, ok := <-resultCh: - // TODO(stream): Check error handling here - if !ok || result == nil { - errCh = nil - // if the resultCh channel is closed, close the responseCh and return - close(responseCh) - return - } - // create a unique message id for each response message which will be used for deduplication - index++ - - // TODO(stream): We need to remove the request message from the tracker once this is completed. - // As we are streaming messages, we need to have some control field to indicate that this is completed - // now, we can do that in the SDK itself. - resp, remove := u.ParseMapResponse(result) - if remove { - u.tracker.RemoveRequest(result.Result.GetUuid()) - - } else if resp != nil { - responseCh <- resp - } - - case err := <-mapErrCh: - // TODO(stream): Check error handling here - // 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) - //} - } - } - }() + //go func() { + // + // index := 0 + // resultCh, mapErrCh := u.client.MapFn(ctx, mapRequestCh) + // for { + // select { + // case result, ok := <-resultCh: + // // TODO(stream): Check error handling here + // if !ok || result == nil { + // errCh = nil + // // if the resultCh channel is closed, close the responseCh and return + // close(responseCh) + // return + // } + // // create a unique message id for each response message which will be used for deduplication + // index++ + // + // // TODO(stream): We need to remove the request message from the tracker once this is completed. + // // As we are streaming messages, we need to have some control field to indicate that this is completed + // // now, we can do that in the SDK itself. + // resp, remove := u.ParseMapResponse(result) + // if remove { + // u.tracker.RemoveRequest(result.Result.GetUuid()) + // + // } else if resp != nil { + // responseCh <- resp + // } + // + // case err := <-mapErrCh: + // // TODO(stream): Check error handling here + // // 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) + // //} + // } + // } + //}() // create ReduceRequest from TimedWindowRequest and send it to reduceRequests channel for AsyncReduceFn go func() { @@ -118,7 +120,7 @@ func (u *GRPCBasedFlatmap) ApplyMap(ctx context.Context, messageStream <-chan *i for { select { case msg, ok := <-messageStream: - //log.Println("MYDEBUG: reading for messages here") + //log.Println("MYDEBUG: reading for messages here ", time.Now()) // if the requestsStream is closed or if the message is nil, return if !ok || msg == nil { //return @@ -126,13 +128,47 @@ func (u *GRPCBasedFlatmap) ApplyMap(ctx context.Context, messageStream <-chan *i d := u.createFlatmapRequest(msg) // send the datum to reduceRequests channel, handle the case when the context is canceled - select { - // TODO(stream): Check the context end here - case mapRequestCh <- d: - //log.Println("MYDEBUG: send the message here", d.Uuid) - //case <-ctx.Done(): - // return + //select { + //// TODO(stream): Check the context end here + //case mapRequestCh <- d: + // //log.Println("MYDEBUG: send the message here", d.Uuid) + // //case <-ctx.Done(): + // // return + //} + + //resp, remove := u.ParseMapResponse(result) + //if remove { + // u.tracker.RemoveRequest(result.Result.GetUuid()) + // + //} else if resp != nil { + // responseCh <- resp + //} + // + + // BYPASS THE UDF BY UNCOMMENTING THIS + taggedMessage := &isb.WriteMessage{ + Message: isb.Message{ + Header: isb.Header{ + MessageInfo: msg.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-%d", msg.ReadOffset.String(), u.idx), + Keys: msg.Keys, + }, + Body: isb.Body{ + Payload: msg.Payload, + }, + }, + Tags: msg.Keys, + } + u.idx += 1 + //u.tracker.IncrementRespIdx(uid) + res := &types.ResponseFlatmap{ + ParentMessage: msg, + Uid: d.Uuid, + RespMessage: taggedMessage, } + responseCh <- res // TODO(stream): Check the context end here, need to invoke shutdown //case <-ctx.Done(): // if the context is done, don't send any more datum to reduceRequests channel // return diff --git a/pkg/flatmap/tracker/tracker.go b/pkg/flatmap/tracker/tracker.go index 24d8b76499..67e635561a 100644 --- a/pkg/flatmap/tracker/tracker.go +++ b/pkg/flatmap/tracker/tracker.go @@ -12,12 +12,15 @@ 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), } } @@ -28,13 +31,13 @@ func GetNewId() string { func (t *Tracker) AddRequest(msg *isb.ReadMessage) string { id := GetNewId() - t.requestMap.Store(id, msg) + //t.requestMap.Store(id, msg) + t.Set(id, msg) return id } func (t *Tracker) GetRequest(id string) (*isb.ReadMessage, bool) { - val, ok := t.requestMap.Load(id) - return val.(*isb.ReadMessage), ok + return t.Get(id) } func (t *Tracker) NewResponse(id string) { @@ -60,6 +63,26 @@ func (t *Tracker) GetIdx(id string) (int, bool) { } func (t *Tracker) RemoveRequest(id string) { - t.requestMap.Delete(id) + //t.requestMap.Delete(id) //t.responseIdx.Delete(id) + t.Delete(id) +} + +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) }