[Go SDK] Timers with new datalayer

sdks/go/examples/timer_wordcap/wordcap.go

@@ -0,0 +1,139 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You 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.
+
+// timer_wordcap is a toy streaming pipeline that demonstrates the use of State and Timers.
+// Periodic Impulse is used as a streaming source that produces sequence of elements upto 5 minutes
+// from the start of the pipeline every 5 seconds. These elements are keyed and fed to the Stateful DoFn
+// where state and timers are set and cleared. Since this pipeline uses a Periodic Impulse,
+// the pipeline is terminated automatically after it is done producing elements for 5 minutes.
+package main
+
+import (
+	"bytes"
+	"context"
+	"encoding/binary"
+	"flag"
+	"fmt"
+	"time"
+
+	"github.com/apache/beam/sdks/v2/go/pkg/beam"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/mtime"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/state"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/timers"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/log"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/register"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/transforms/periodic"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/x/beamx"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/x/debug"
+)
+
+type Stateful struct {
+	ElementBag state.Bag[string]
+	TimerTime  state.Value[int64]
+	MinTime    state.Combining[int64, int64, int64]
+
+	OutputState timers.ProcessingTime
+}
+
+func NewStateful() *Stateful {
+	return &Stateful{
+		ElementBag: state.MakeBagState[string]("elementBag"),
+		TimerTime:  state.MakeValueState[int64]("timerTime"),
+		MinTime: state.MakeCombiningState[int64, int64, int64]("minTiInBag", func(a, b int64) int64 {
+			if a < b {
+				return a
+			}
+			return b
+		}),
+
+		OutputState: timers.InProcessingTime("outputState"),
+	}
+}
+
+func (s *Stateful) OnTimer(ctx context.Context, ts beam.EventTime, tp timers.Provider, key, timerKey, timerTag string, emit func(string, string)) {
+	switch timerKey {
+	case "outputState":
+		log.Infof(ctx, "Timer outputState fired on stateful for element: %v.", key)
+		s.OutputState.Set(tp, ts.ToTime().Add(5*time.Second), timers.WithTag("1"))
+		switch timerTag {
+		case "1":
+			s.OutputState.Clear(tp)
+			log.Infof(ctx, "Timer with tag 1 fired on outputState stateful DoFn.")
+			emit(timerKey, timerTag)
+		}
+	}
+}
+
+func (s *Stateful) ProcessElement(ctx context.Context, ts beam.EventTime, sp state.Provider, tp timers.Provider, key, word string, emit func(string, string)) error {
+	s.ElementBag.Add(sp, word)
+	s.MinTime.Add(sp, int64(ts))
+
+	toFire, ok, err := s.TimerTime.Read(sp)
+	if err != nil {
+		return err
+	}
+	if !ok {
+		toFire = int64(mtime.Now().Add(1 * time.Minute))
+	}
+	minTime, _, err := s.MinTime.Read(sp)
+	if err != nil {
+		return err
+	}
+
+	s.OutputState.Set(tp, time.UnixMilli(toFire), timers.WithOutputTimestamp(time.UnixMilli(minTime)), timers.WithTag(word))
+	s.TimerTime.Write(sp, toFire)
+
+	return nil
+}
+
+func init() {
+	register.DoFn7x1[context.Context, beam.EventTime, state.Provider, timers.Provider, string, string, func(string, string), error](&Stateful{})
+	register.Emitter2[string, string]()
+	register.Emitter2[beam.EventTime, int64]()
+}
+
+func main() {
+	flag.Parse()
+	beam.Init()
+
+	ctx := context.Background()
+
+	p := beam.NewPipeline()
+	s := p.Root()
+
+	out := periodic.Impulse(s, time.Now(), time.Now().Add(5*time.Minute), 5*time.Second, true)
+
+	intOut := beam.ParDo(s, func(b []byte) int64 {
+		var val int64
+		buf := bytes.NewReader(b)
+		binary.Read(buf, binary.BigEndian, &val)
+		return val
+	}, out)
+
+	str := beam.ParDo(s, func(b int64) string {
+		return fmt.Sprintf("%03d", b)
+	}, intOut)
+
+	keyed := beam.ParDo(s, func(ctx context.Context, ts beam.EventTime, s string) (string, string) {
+		return "test", s
+	}, str)
+
+	timed := beam.ParDo(s, NewStateful(), keyed)
+	debug.Printf(s, "post stateful: %v", timed)
+
+	if err := beamx.Run(context.Background(), p); err != nil {
+		log.Exitf(ctx, "Failed to execute job: %v", err)
+	}
+}

sdks/go/pkg/beam/core/funcx/fn.go

@@ -21,6 +21,7 @@ import (
 
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/sdf"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/state"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/timers"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/reflectx"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors"
@@ -85,6 +86,8 @@ const (
 	FnWatermarkEstimator FnParamKind = 0x1000
 	// FnStateProvider indicates a function input parameter that implements state.Provider
 	FnStateProvider FnParamKind = 0x2000
+	// FnTimerProvider indicates a function input parameter that implements timer.Provider
+	FnTimerProvider FnParamKind = 0x4000
 )
 
 func (k FnParamKind) String() string {
@@ -117,6 +120,8 @@ func (k FnParamKind) String() string {
 		return "WatermarkEstimator"
 	case FnStateProvider:
 		return "StateProvider"
+	case FnTimerProvider:
+		return "TimerProvider"
 	default:
 		return fmt.Sprintf("%v", int(k))
 	}
@@ -305,6 +310,17 @@ func (u *Fn) StateProvider() (pos int, exists bool) {
 	return -1, false
 }
 
+// TimerProvider returns (index, true) iff the function expects a
+// parameter that implements timers.Provider.
+func (u *Fn) TimerProvider() (pos int, exists bool) {
+	for i, p := range u.Param {
+		if p.Kind == FnTimerProvider {
+			return i, true
+		}
+	}
+	return -1, false
+}
+
 // WatermarkEstimator returns (index, true) iff the function expects a
 // parameter that implements sdf.WatermarkEstimator.
 func (u *Fn) WatermarkEstimator() (pos int, exists bool) {
@@ -392,6 +408,8 @@ func New(fn reflectx.Func) (*Fn, error) {
 			kind = FnBundleFinalization
 		case t == state.ProviderType:
 			kind = FnStateProvider
+		case t == timers.ProviderType:
+			kind = FnTimerProvider
 		case t == reflectx.Type:
 			kind = FnType
 		case t.Implements(reflect.TypeOf((*sdf.RTracker)(nil)).Elem()):
@@ -482,7 +500,7 @@ func SubReturns(list []ReturnParam, indices ...int) []ReturnParam {
 }
 
 // The order of present parameters and return values must be as follows:
-// func(FnContext?, FnPane?, FnWindow?, FnEventTime?, FnWatermarkEstimator?, FnType?, FnBundleFinalization?, FnRTracker?, FnStateProvider?, (FnValue, SideInput*)?, FnEmit*) (RetEventTime?, RetOutput?, RetError?)
+// func(FnContext?, FnPane?, FnWindow?, FnEventTime?, FnWatermarkEstimator?, FnType?, FnBundleFinalization?, FnRTracker?, FnStateProvider?, FnTimerProvider?, (FnValue, SideInput*)?, FnEmit*) (RetEventTime?, RetOutput?, RetError?)
 //
 //	where ? indicates 0 or 1, and * indicates any number.
 //	and  a SideInput is one of FnValue or FnIter or FnReIter
@@ -517,6 +535,7 @@ var (
 	errRTrackerPrecedence                = errors.New("may only have a single sdf.RTracker parameter and it must precede the main input parameter")
 	errBundleFinalizationPrecedence      = errors.New("may only have a single BundleFinalization parameter and it must precede the main input parameter")
 	errStateProviderPrecedence           = errors.New("may only have a single state.Provider parameter and it must precede the main input parameter")
+	errTimerProviderPrecedence           = errors.New("may only have a single timer.Provider parameter and it must precede the main input parameter")
 	errInputPrecedence                   = errors.New("inputs parameters must precede emit function parameters")
 )
 
@@ -535,6 +554,7 @@ const (
 	psRTracker
 	psBundleFinalization
 	psStateProvider
+	psTimerProvider
 )
 
 func nextParamState(cur paramState, transition FnParamKind) (paramState, error) {
@@ -559,6 +579,8 @@ func nextParamState(cur paramState, transition FnParamKind) (paramState, error)
 			return psRTracker, nil
 		case FnStateProvider:
 			return psStateProvider, nil
+		case FnTimerProvider:
+			return psTimerProvider, nil
 		}
 	case psContext:
 		switch transition {
@@ -578,6 +600,8 @@ func nextParamState(cur paramState, transition FnParamKind) (paramState, error)
 			return psRTracker, nil
 		case FnStateProvider:
 			return psStateProvider, nil
+		case FnTimerProvider:
+			return psTimerProvider, nil
 		}
 	case psPane:
 		switch transition {
@@ -595,6 +619,8 @@ func nextParamState(cur paramState, transition FnParamKind) (paramState, error)
 			return psRTracker, nil
 		case FnStateProvider:
 			return psStateProvider, nil
+		case FnTimerProvider:
+			return psTimerProvider, nil
 		}
 	case psWindow:
 		switch transition {
@@ -610,6 +636,8 @@ func nextParamState(cur paramState, transition FnParamKind) (paramState, error)
 			return psRTracker, nil
 		case FnStateProvider:
 			return psStateProvider, nil
+		case FnTimerProvider:
+			return psTimerProvider, nil
 		}
 	case psEventTime:
 		switch transition {
@@ -623,6 +651,8 @@ func nextParamState(cur paramState, transition FnParamKind) (paramState, error)
 			return psRTracker, nil
 		case FnStateProvider:
 			return psStateProvider, nil
+		case FnTimerProvider:
+			return psTimerProvider, nil
 		}
 	case psWatermarkEstimator:
 		switch transition {
@@ -634,6 +664,8 @@ func nextParamState(cur paramState, transition FnParamKind) (paramState, error)
 			return psRTracker, nil
 		case FnStateProvider:
 			return psStateProvider, nil
+		case FnTimerProvider:
+			return psTimerProvider, nil
 		}
 	case psType:
 		switch transition {
@@ -643,20 +675,31 @@ func nextParamState(cur paramState, transition FnParamKind) (paramState, error)
 			return psRTracker, nil
 		case FnStateProvider:
 			return psStateProvider, nil
+		case FnTimerProvider:
+			return psTimerProvider, nil
 		}
 	case psBundleFinalization:
 		switch transition {
 		case FnRTracker:
 			return psRTracker, nil
 		case FnStateProvider:
 			return psStateProvider, nil
+		case FnTimerProvider:
+			return psTimerProvider, nil
 		}
 	case psRTracker:
 		switch transition {
 		case FnStateProvider:
 			return psStateProvider, nil
+		case FnTimerProvider:
+			return psTimerProvider, nil
 		}
 	case psStateProvider:
+		switch transition {
+		case FnTimerProvider:
+			return psTimerProvider, nil
+		}
+	case psTimerProvider:
 		// Completely handled by the default clause
 	case psInput:
 		switch transition {
@@ -689,6 +732,8 @@ func nextParamState(cur paramState, transition FnParamKind) (paramState, error)
 		return -1, errRTrackerPrecedence
 	case FnStateProvider:
 		return -1, errStateProviderPrecedence
+	case FnTimerProvider:
+		return -1, errTimerProviderPrecedence
 	case FnIter, FnReIter, FnValue, FnMultiMap:
 		return psInput, nil
 	case FnEmit: