github.com/apache/beam/sdks/v2@v2.48.2/go/examples/timer_wordcap/wordcap.go (about)

     1  // Licensed to the Apache Software Foundation (ASF) under one or more
     2  // contributor license agreements.  See the NOTICE file distributed with
     3  // this work for additional information regarding copyright ownership.
     4  // The ASF licenses this file to You under the Apache License, Version 2.0
     5  // (the "License"); you may not use this file except in compliance with
     6  // the License.  You may obtain a copy of the License at
     7  //
     8  //    http://www.apache.org/licenses/LICENSE-2.0
     9  //
    10  // Unless required by applicable law or agreed to in writing, software
    11  // distributed under the License is distributed on an "AS IS" BASIS,
    12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  // See the License for the specific language governing permissions and
    14  // limitations under the License.
    15  
    16  // timer_wordcap is a toy streaming pipeline that demonstrates the use of State and Timers.
    17  // Periodic Impulse is used as a streaming source that produces sequence of elements upto 5 minutes
    18  // from the start of the pipeline every 5 seconds. These elements are keyed and fed to the Stateful DoFn
    19  // where state and timers are set and cleared. Since this pipeline uses a Periodic Impulse,
    20  // the pipeline is terminated automatically after it is done producing elements for 5 minutes.
    21  package main
    22  
    23  import (
    24  	"bytes"
    25  	"context"
    26  	"encoding/binary"
    27  	"flag"
    28  	"fmt"
    29  	"time"
    30  
    31  	"github.com/apache/beam/sdks/v2/go/pkg/beam"
    32  	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/mtime"
    33  	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/state"
    34  	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/timers"
    35  	"github.com/apache/beam/sdks/v2/go/pkg/beam/log"
    36  	"github.com/apache/beam/sdks/v2/go/pkg/beam/register"
    37  	"github.com/apache/beam/sdks/v2/go/pkg/beam/transforms/periodic"
    38  	"github.com/apache/beam/sdks/v2/go/pkg/beam/x/beamx"
    39  	"github.com/apache/beam/sdks/v2/go/pkg/beam/x/debug"
    40  )
    41  
    42  type Stateful struct {
    43  	ElementBag state.Bag[string]
    44  	TimerTime  state.Value[int64]
    45  	MinTime    state.Combining[int64, int64, int64]
    46  
    47  	OutputState timers.ProcessingTime
    48  }
    49  
    50  func NewStateful() *Stateful {
    51  	return &Stateful{
    52  		ElementBag: state.MakeBagState[string]("elementBag"),
    53  		TimerTime:  state.MakeValueState[int64]("timerTime"),
    54  		MinTime: state.MakeCombiningState[int64, int64, int64]("minTiInBag", func(a, b int64) int64 {
    55  			if a < b {
    56  				return a
    57  			}
    58  			return b
    59  		}),
    60  
    61  		OutputState: timers.InProcessingTime("outputState"),
    62  	}
    63  }
    64  
    65  func (s *Stateful) OnTimer(ctx context.Context, ts beam.EventTime, tp timers.Provider, key, timerKey, timerTag string, emit func(string, string)) {
    66  	switch timerKey {
    67  	case "outputState":
    68  		log.Infof(ctx, "Timer outputState fired on stateful for element: %v.", key)
    69  		s.OutputState.Set(tp, ts.ToTime().Add(5*time.Second), timers.WithTag("1"))
    70  		switch timerTag {
    71  		case "1":
    72  			s.OutputState.Clear(tp)
    73  			log.Infof(ctx, "Timer with tag 1 fired on outputState stateful DoFn.")
    74  			emit(timerKey, timerTag)
    75  		}
    76  	}
    77  }
    78  
    79  func (s *Stateful) ProcessElement(ctx context.Context, ts beam.EventTime, sp state.Provider, tp timers.Provider, key, word string, emit func(string, string)) error {
    80  	s.ElementBag.Add(sp, word)
    81  	s.MinTime.Add(sp, int64(ts))
    82  
    83  	toFire, ok, err := s.TimerTime.Read(sp)
    84  	if err != nil {
    85  		return err
    86  	}
    87  	if !ok {
    88  		toFire = int64(mtime.Now().Add(1 * time.Minute))
    89  	}
    90  	minTime, _, err := s.MinTime.Read(sp)
    91  	if err != nil {
    92  		return err
    93  	}
    94  
    95  	s.OutputState.Set(tp, time.UnixMilli(toFire), timers.WithOutputTimestamp(time.UnixMilli(minTime)), timers.WithTag(word))
    96  	s.TimerTime.Write(sp, toFire)
    97  
    98  	return nil
    99  }
   100  
   101  func init() {
   102  	register.DoFn7x1[context.Context, beam.EventTime, state.Provider, timers.Provider, string, string, func(string, string), error](&Stateful{})
   103  	register.Emitter2[string, string]()
   104  	register.Emitter2[beam.EventTime, int64]()
   105  }
   106  
   107  func main() {
   108  	flag.Parse()
   109  	beam.Init()
   110  
   111  	ctx := context.Background()
   112  
   113  	p := beam.NewPipeline()
   114  	s := p.Root()
   115  
   116  	out := periodic.Impulse(s, time.Now(), time.Now().Add(5*time.Minute), 5*time.Second, true)
   117  
   118  	intOut := beam.ParDo(s, func(b []byte) int64 {
   119  		var val int64
   120  		buf := bytes.NewReader(b)
   121  		binary.Read(buf, binary.BigEndian, &val)
   122  		return val
   123  	}, out)
   124  
   125  	str := beam.ParDo(s, func(b int64) string {
   126  		return fmt.Sprintf("%03d", b)
   127  	}, intOut)
   128  
   129  	keyed := beam.ParDo(s, func(ctx context.Context, ts beam.EventTime, s string) (string, string) {
   130  		return "test", s
   131  	}, str)
   132  
   133  	timed := beam.ParDo(s, NewStateful(), keyed)
   134  	debug.Printf(s, "post stateful: %v", timed)
   135  
   136  	if err := beamx.Run(context.Background(), p); err != nil {
   137  		log.Exitf(ctx, "Failed to execute job: %v", err)
   138  	}
   139  }