github.com/apache/beam/sdks/v2@v2.48.2/go/examples/native_wordcap/nativepubsubio/subscriptiontracker.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 nativepubsubio
    17  
    18  // The SubscriptionRTracker maintains a single entry string to keep up with
    19  // the PubSub subscription being used in the NativeRead SDF.
    20  type SubscriptionRTracker struct {
    21  	Subscription string
    22  	Done         bool
    23  }
    24  
    25  // NewSubscriptionRTracker returns an RTracker wrapping the
    26  // provided subscription and a "Done" boolean.
    27  func NewSubscriptionRTracker(entry string) *SubscriptionRTracker {
    28  	return &SubscriptionRTracker{Subscription: entry, Done: false}
    29  }
    30  
    31  // TryClaim returns true iff the given position is a string and matches the underlying
    32  // subscription ID.
    33  func (s *SubscriptionRTracker) TryClaim(pos any) bool {
    34  	posString, ok := pos.(string)
    35  	return ok && posString == s.Subscription
    36  }
    37  
    38  // TrySplit is a no-op for the StaticRTracker in the normal case and moves the subscription
    39  // to the residual in the checkpointing case, marking itself as done to keep the logical checks
    40  // around SDF data loss happy.
    41  func (s *SubscriptionRTracker) TrySplit(frac float64) (primary, residual any, err error) {
    42  	if frac == 0.0 {
    43  		resid := s.Subscription
    44  		s.Subscription = ""
    45  		s.Done = true
    46  		return "", resid, nil
    47  	}
    48  	return s.Subscription, "", nil
    49  }
    50  
    51  // GetError is a no-op.
    52  func (s *SubscriptionRTracker) GetError() error {
    53  	return nil
    54  }
    55  
    56  // GetProgress returns complete just so the runner doesn't try to do much in the way of
    57  // splitting.
    58  func (s *SubscriptionRTracker) GetProgress() (done float64, remaining float64) {
    59  	done = 1.0
    60  	remaining = 0.0
    61  	return
    62  }
    63  
    64  // IsDone returns whether or not the StaticRTracker is complete (e.g. has stopped processing.)
    65  func (s *SubscriptionRTracker) IsDone() bool {
    66  	return s.Done
    67  }
    68  
    69  // IsBounded always returns false, as the StaticRTracker represents an unbounded number
    70  // of reads from PubSub.
    71  func (s *SubscriptionRTracker) IsBounded() bool {
    72  	return false
    73  }
    74  
    75  // GetRestriction returns the name of the subscription.
    76  func (s *SubscriptionRTracker) GetRestriction() any {
    77  	return s.Subscription
    78  }