github.com/apache/beam/sdks/v2@v2.48.2/go/test/integration/primitives/drain.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  	"math"
    21  	"time"
    22  
    23  	"github.com/apache/beam/sdks/v2/go/pkg/beam"
    24  	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/sdf"
    25  	"github.com/apache/beam/sdks/v2/go/pkg/beam/io/rtrackers/offsetrange"
    26  	"github.com/apache/beam/sdks/v2/go/pkg/beam/log"
    27  	"github.com/apache/beam/sdks/v2/go/pkg/beam/register"
    28  )
    29  
    30  func init() {
    31  	register.DoFn4x1[context.Context, *sdf.LockRTracker, []byte, func(int64), sdf.ProcessContinuation](&TruncateFn{})
    32  
    33  	register.Emitter1[int64]()
    34  }
    35  
    36  // RangeEstimator implements the offsetrange.RangeEndEstimator interface.
    37  // It provides the estimated end for a restriction.
    38  type RangeEstimator struct {
    39  	end int64
    40  }
    41  
    42  // Estimate returns the estimated end.
    43  func (r *RangeEstimator) Estimate() int64 {
    44  	return r.end
    45  }
    46  
    47  // SetEstimate sets the estimated end.
    48  func (r *RangeEstimator) SetEstimate(estimate int64) {
    49  	r.end = estimate
    50  }
    51  
    52  // TruncateFn is an SDF.
    53  type TruncateFn struct {
    54  	Estimator RangeEstimator
    55  }
    56  
    57  // CreateInitialRestriction creates an initial restriction
    58  func (fn *TruncateFn) CreateInitialRestriction(_ []byte) offsetrange.Restriction {
    59  	return offsetrange.Restriction{
    60  		Start: int64(1),
    61  		End:   int64(math.MaxInt64),
    62  	}
    63  }
    64  
    65  // CreateTracker wraps the given restriction into a LockRTracker type.
    66  func (fn *TruncateFn) CreateTracker(rest offsetrange.Restriction) *sdf.LockRTracker {
    67  	fn.Estimator = RangeEstimator{int64(10)}
    68  	tracker, err := offsetrange.NewGrowableTracker(rest, &fn.Estimator)
    69  	if err != nil {
    70  		panic(err)
    71  	}
    72  	return sdf.NewLockRTracker(tracker)
    73  }
    74  
    75  // RestrictionSize returns the size of the current restriction
    76  func (fn *TruncateFn) RestrictionSize(_ []byte, rest offsetrange.Restriction) float64 {
    77  	return rest.Size()
    78  }
    79  
    80  // SplitRestriction is similar to the one used in checkpointing.go test.
    81  func (fn *TruncateFn) SplitRestriction(_ []byte, rest offsetrange.Restriction) []offsetrange.Restriction {
    82  	return rest.EvenSplits(2)
    83  }
    84  
    85  // TruncateRestriction truncates the restriction during drain.
    86  func (fn *TruncateFn) TruncateRestriction(ctx context.Context, rt *sdf.LockRTracker, _ []byte) offsetrange.Restriction {
    87  	rest := rt.GetRestriction().(offsetrange.Restriction)
    88  	start := rest.Start
    89  	newEnd := start + 20
    90  
    91  	done, remaining := rt.GetProgress()
    92  	log.Infof(ctx, "Draining at: done %v, remaining %v, start %v, end %v, newEnd %v", done, remaining, start, rest.End, newEnd)
    93  
    94  	return offsetrange.Restriction{
    95  		Start: start,
    96  		End:   newEnd,
    97  	}
    98  }
    99  
   100  // ProcessElement continually gets the start position of the restriction and emits the element as it is.
   101  func (fn *TruncateFn) ProcessElement(ctx context.Context, rt *sdf.LockRTracker, _ []byte, emit func(int64)) sdf.ProcessContinuation {
   102  	position := rt.GetRestriction().(offsetrange.Restriction).Start
   103  	for {
   104  		if rt.TryClaim(position) {
   105  			log.Infof(ctx, "Claimed position: %v", position)
   106  			// Successful claim, emit the value and move on.
   107  			emit(position)
   108  			position++
   109  		} else if rt.GetError() != nil || rt.IsDone() {
   110  			// Stop processing on error or completion
   111  			if err := rt.GetError(); err != nil {
   112  				log.Errorf(ctx, "error in restriction tracker, got %v", err)
   113  			}
   114  			log.Infof(ctx, "Restriction done at position %v.", position)
   115  			return sdf.StopProcessing()
   116  		} else {
   117  			log.Infof(ctx, "Checkpointed at position %v, resuming later.", position)
   118  			// Resume later.
   119  			return sdf.ResumeProcessingIn(5 * time.Second)
   120  		}
   121  		time.Sleep(1 * time.Second)
   122  	}
   123  }
   124  
   125  // Drain tests the SDF truncation during drain.
   126  func Drain(s beam.Scope) {
   127  	beam.Init()
   128  	s.Scope("truncate")
   129  	beam.ParDo(s, &TruncateFn{}, beam.Impulse(s))
   130  }