github.com/apache/beam/sdks/v2@v2.48.2/go/test/integration/primitives/checkpointing.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  package primitives
    17  
    18  import (
    19  	"context"
    20  	"time"
    21  
    22  	"github.com/apache/beam/sdks/v2/go/pkg/beam"
    23  	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/sdf"
    24  	"github.com/apache/beam/sdks/v2/go/pkg/beam/io/rtrackers/offsetrange"
    25  	"github.com/apache/beam/sdks/v2/go/pkg/beam/log"
    26  	"github.com/apache/beam/sdks/v2/go/pkg/beam/register"
    27  	"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/passert"
    28  )
    29  
    30  func init() {
    31  	register.DoFn3x1[*sdf.LockRTracker, []byte, func(int64), sdf.ProcessContinuation](&selfCheckpointingDoFn{})
    32  	register.Emitter1[int64]()
    33  }
    34  
    35  type selfCheckpointingDoFn struct{}
    36  
    37  // CreateInitialRestriction creates the restriction being used by the SDF. In this case, the range
    38  // of values produced by the restriction is [Start, End).
    39  func (fn *selfCheckpointingDoFn) CreateInitialRestriction(_ []byte) offsetrange.Restriction {
    40  	return offsetrange.Restriction{
    41  		Start: int64(0),
    42  		End:   int64(10),
    43  	}
    44  }
    45  
    46  // CreateTracker wraps the given restriction into a LockRTracker type.
    47  func (fn *selfCheckpointingDoFn) CreateTracker(rest offsetrange.Restriction) *sdf.LockRTracker {
    48  	return sdf.NewLockRTracker(offsetrange.NewTracker(rest))
    49  }
    50  
    51  // RestrictionSize returns the size of the current restriction
    52  func (fn *selfCheckpointingDoFn) RestrictionSize(_ []byte, rest offsetrange.Restriction) float64 {
    53  	return rest.Size()
    54  }
    55  
    56  // SplitRestriction modifies the offsetrange.Restriction's sized restriction function to produce a size-zero restriction
    57  // at the end of execution.
    58  func (fn *selfCheckpointingDoFn) SplitRestriction(_ []byte, rest offsetrange.Restriction) []offsetrange.Restriction {
    59  	size := int64(10)
    60  	s := rest.Start
    61  	var splits []offsetrange.Restriction
    62  	for e := s + size; e <= rest.End; s, e = e, e+size {
    63  		splits = append(splits, offsetrange.Restriction{Start: s, End: e})
    64  	}
    65  	splits = append(splits, offsetrange.Restriction{Start: s, End: rest.End})
    66  	return splits
    67  }
    68  
    69  // ProcessElement continually gets the start position of the restriction and emits it as an int64 value before checkpointing.
    70  // This causes the restriction to be split after the claimed work and produce no primary roots.
    71  func (fn *selfCheckpointingDoFn) ProcessElement(rt *sdf.LockRTracker, _ []byte, emit func(int64)) sdf.ProcessContinuation {
    72  	position := rt.GetRestriction().(offsetrange.Restriction).Start
    73  
    74  	counter := 0
    75  	for {
    76  		if rt.TryClaim(position) {
    77  			// Successful claim, emit the value and move on.
    78  			emit(position)
    79  			position++
    80  			counter++
    81  		} else if rt.GetError() != nil || rt.IsDone() {
    82  			// Stop processing on error or completion
    83  			if err := rt.GetError(); err != nil {
    84  				log.Errorf(context.Background(), "error in restriction tracker, got %v", err)
    85  			}
    86  			return sdf.StopProcessing()
    87  		} else {
    88  			// Resume later.
    89  			return sdf.ResumeProcessingIn(5 * time.Second)
    90  		}
    91  
    92  		if counter >= 10 {
    93  			return sdf.ResumeProcessingIn(1 * time.Second)
    94  		}
    95  	}
    96  }
    97  
    98  // Checkpoints is a small test pipeline to establish the correctness of the simple test case.
    99  func Checkpoints(s beam.Scope) {
   100  	s.Scope("checkpoint")
   101  	out := beam.ParDo(s, &selfCheckpointingDoFn{}, beam.Impulse(s))
   102  	passert.Count(s, out, "num ints", 10)
   103  }