github.com/google/fleetspeak@v0.1.15-0.20240426164851-4f31f62c1aea/fleetspeak/src/client/internal/message/sort.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
    16  
    17  import (
    18  	"time"
    19  
    20  	"github.com/google/fleetspeak/fleetspeak/src/client/comms"
    21  	"github.com/google/fleetspeak/fleetspeak/src/client/flow"
    22  
    23  	fspb "github.com/google/fleetspeak/fleetspeak/src/common/proto/fleetspeak"
    24  )
    25  
    26  // SortLoop connects in and out in a sorted manner. That is, it acts essentially
    27  // as a buffered channel between in and out which sorts any messages within
    28  // it. Caller is responsible for implementing any needed size limit. Returns
    29  // when "in" is closed.
    30  func SortLoop(in <-chan comms.MessageInfo, out chan<- comms.MessageInfo, f *flow.Filter) {
    31  	// Keep a slice of messages for each priority level. These are used as fifo
    32  	// queues, appending to the end and retreiving from head.
    33  	var low, medium, high []comms.MessageInfo
    34  
    35  	// Append a message to the correct list.
    36  	appendMI := func(mi comms.MessageInfo) {
    37  		switch mi.M.Priority {
    38  		case fspb.Message_LOW:
    39  			low = append(low, mi)
    40  		case fspb.Message_MEDIUM:
    41  			medium = append(medium, mi)
    42  		case fspb.Message_HIGH:
    43  			high = append(high, mi)
    44  		default:
    45  			medium = append(medium, mi)
    46  		}
    47  	}
    48  
    49  	// Used to poll f occasionally, in case it changes while we are waiting for
    50  	// something.
    51  	filterPoll := time.NewTicker(time.Second)
    52  	defer filterPoll.Stop()
    53  
    54  	for {
    55  		// Try to send from the highest priority non-empty queue. Simultaneously try
    56  		// to receive. (First true case matches, of course.)
    57  		//
    58  		// Don't send from queues which are filtered.
    59  		lf, mf, hf := f.Get()
    60  		switch {
    61  		case len(high) != 0 && !hf:
    62  			select {
    63  			case mi, ok := <-in:
    64  				if !ok {
    65  					return
    66  				}
    67  				appendMI(mi)
    68  			case out <- high[0]:
    69  				high = high[1:]
    70  			case <-filterPoll.C:
    71  			}
    72  		case len(medium) != 0 && !mf:
    73  			select {
    74  			case mi, ok := <-in:
    75  				if !ok {
    76  					return
    77  				}
    78  				appendMI(mi)
    79  			case out <- medium[0]:
    80  				medium = medium[1:]
    81  			case <-filterPoll.C:
    82  			}
    83  		case len(low) != 0 && !lf:
    84  			select {
    85  			case mi, ok := <-in:
    86  				if !ok {
    87  					return
    88  				}
    89  				appendMI(mi)
    90  			case out <- low[0]:
    91  				low = low[1:]
    92  			case <-filterPoll.C:
    93  			}
    94  		default:
    95  			// Nothing to send, we can only receive.
    96  			select {
    97  			case mi, ok := <-in:
    98  				if !ok {
    99  					return
   100  				}
   101  				appendMI(mi)
   102  			case <-filterPoll.C:
   103  			}
   104  		}
   105  	}
   106  }