github.com/google/fleetspeak@v0.1.15-0.20240426164851-4f31f62c1aea/fleetspeak/src/client/internal/message/retry.go (about)

     1  // Copyright 2017 Google Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     https://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Package message implements utility structures and methods used by the FS
    16  // client to manage messages.
    17  package message
    18  
    19  import (
    20  	"google.golang.org/protobuf/proto"
    21  
    22  	"github.com/google/fleetspeak/fleetspeak/src/client/comms"
    23  	"github.com/google/fleetspeak/fleetspeak/src/client/service"
    24  	"github.com/google/fleetspeak/fleetspeak/src/client/stats"
    25  )
    26  
    27  type sizedMessage struct {
    28  	size int
    29  	m    service.AckMessage
    30  }
    31  
    32  // RetryLoop is a loop which reads from in and writes to out.
    33  //
    34  // Messages are considered pending until MessageInfo.Ack is called and output
    35  // again if MessageInfo.Nack is called. The loop will pause reading from input
    36  // when at least maxSize bytes of messages, or maxCount messages are pending.
    37  //
    38  // To shutdown gracefully, close out and Ack any pending messages.
    39  func RetryLoop(in <-chan service.AckMessage, out chan<- comms.MessageInfo, stats stats.RetryLoopCollector, maxSize, maxCount int) {
    40  	// Used to send acks/nacks back to this loop. Buffered to prevent ack,nack
    41  	// callbacks from ever blocking.
    42  	acks := make(chan sizedMessage, maxCount)
    43  	nacks := make(chan sizedMessage, maxCount)
    44  
    45  	makeInfo := func(sm sizedMessage) comms.MessageInfo {
    46  		return comms.MessageInfo{
    47  			M: sm.m.M,
    48  			Ack: func() {
    49  				if sm.m.Ack != nil {
    50  					sm.m.Ack()
    51  				}
    52  				acks <- sm
    53  			},
    54  			Nack: func() { nacks <- sm },
    55  		}
    56  	}
    57  
    58  	var size, count int
    59  	var optIn <-chan service.AckMessage
    60  	for {
    61  		if size >= maxSize || count >= maxCount {
    62  			optIn = nil
    63  		} else {
    64  			optIn = in
    65  		}
    66  
    67  		select {
    68  		case sm := <-acks:
    69  			size -= sm.size
    70  			count--
    71  			stats.MessageAcknowledged(sm.m.M, sm.size)
    72  		case sm := <-nacks:
    73  			out <- makeInfo(sm)
    74  			stats.BeforeMessageRetry(sm.m.M)
    75  		case m, ok := <-optIn:
    76  			if !ok {
    77  				return
    78  			}
    79  			sm := sizedMessage{proto.Size(m.M), m}
    80  			size += sm.size
    81  			count++
    82  			stats.MessagePending(m.M, sm.size)
    83  			out <- makeInfo(sm)
    84  		}
    85  	}
    86  }