github.com/apache/beam/sdks/v2@v2.48.2/go/examples/snippets/12splittabledofns.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 snippets
    17  
    18  import (
    19  	"time"
    20  
    21  	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/sdf"
    22  	"github.com/apache/beam/sdks/v2/go/pkg/beam/io/rtrackers/offsetrange"
    23  )
    24  
    25  type Record struct{}
    26  
    27  type SomeService struct {
    28  	ThrottlingErr error
    29  }
    30  
    31  func (s *SomeService) readNextRecords(position any) ([]Record, error) {
    32  	return []Record{}, nil
    33  }
    34  
    35  type checkpointingSplittableDoFn struct {
    36  	ExternalService SomeService
    37  }
    38  
    39  // [START self_checkpoint]
    40  func (fn *checkpointingSplittableDoFn) ProcessElement(rt *sdf.LockRTracker, emit func(Record)) (sdf.ProcessContinuation, error) {
    41  	position := rt.GetRestriction().(offsetrange.Restriction).Start
    42  	for {
    43  		records, err := fn.ExternalService.readNextRecords(position)
    44  
    45  		if err != nil {
    46  			if err == fn.ExternalService.ThrottlingErr {
    47  				// Resume at a later time to avoid throttling.
    48  				return sdf.ResumeProcessingIn(60 * time.Second), nil
    49  			}
    50  			return sdf.StopProcessing(), err
    51  		}
    52  
    53  		if len(records) == 0 {
    54  			// Wait for data to be available.
    55  			return sdf.ResumeProcessingIn(10 * time.Second), nil
    56  		}
    57  		for _, record := range records {
    58  			if !rt.TryClaim(position) {
    59  				// Records have been claimed, finish processing.
    60  				return sdf.StopProcessing(), nil
    61  			}
    62  			position += 1
    63  
    64  			emit(record)
    65  		}
    66  	}
    67  }
    68  
    69  // [END self_checkpoint]