istio.io/istio@v0.0.0-20240520182934-d79c90f27776/pilot/pkg/xds/pushqueue.go (about)

     1  // Copyright Istio Authors
     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  //     http://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 xds
    16  
    17  import (
    18  	"sync"
    19  
    20  	"istio.io/istio/pilot/pkg/model"
    21  )
    22  
    23  type PushQueue struct {
    24  	cond *sync.Cond
    25  
    26  	// pending stores all connections in the queue. If the same connection is enqueued again,
    27  	// the PushRequest will be merged.
    28  	pending map[*Connection]*model.PushRequest
    29  
    30  	// queue maintains ordering of the queue
    31  	queue []*Connection
    32  
    33  	// processing stores all connections that have been Dequeue(), but not MarkDone().
    34  	// The value stored will be initially be nil, but may be populated if the connection is Enqueue().
    35  	// If model.PushRequest is not nil, it will be Enqueued again once MarkDone has been called.
    36  	processing map[*Connection]*model.PushRequest
    37  
    38  	shuttingDown bool
    39  }
    40  
    41  func NewPushQueue() *PushQueue {
    42  	return &PushQueue{
    43  		pending:    make(map[*Connection]*model.PushRequest),
    44  		processing: make(map[*Connection]*model.PushRequest),
    45  		cond:       sync.NewCond(&sync.Mutex{}),
    46  	}
    47  }
    48  
    49  // Enqueue will mark a proxy as pending a push. If it is already pending, pushInfo will be merged.
    50  // ServiceEntry updates will be added together, and full will be set if either were full
    51  func (p *PushQueue) Enqueue(con *Connection, pushRequest *model.PushRequest) {
    52  	p.cond.L.Lock()
    53  	defer p.cond.L.Unlock()
    54  
    55  	if p.shuttingDown {
    56  		return
    57  	}
    58  
    59  	// If its already in progress, merge the info and return
    60  	if request, f := p.processing[con]; f {
    61  		p.processing[con] = request.CopyMerge(pushRequest)
    62  		return
    63  	}
    64  
    65  	if request, f := p.pending[con]; f {
    66  		p.pending[con] = request.CopyMerge(pushRequest)
    67  		return
    68  	}
    69  
    70  	p.pending[con] = pushRequest
    71  	p.queue = append(p.queue, con)
    72  	// Signal waiters on Dequeue that a new item is available
    73  	p.cond.Signal()
    74  }
    75  
    76  // Remove a proxy from the queue. If there are no proxies ready to be removed, this will block
    77  func (p *PushQueue) Dequeue() (con *Connection, request *model.PushRequest, shutdown bool) {
    78  	p.cond.L.Lock()
    79  	defer p.cond.L.Unlock()
    80  
    81  	// Block until there is one to remove. Enqueue will signal when one is added.
    82  	for len(p.queue) == 0 && !p.shuttingDown {
    83  		p.cond.Wait()
    84  	}
    85  
    86  	if len(p.queue) == 0 {
    87  		// We must be shutting down.
    88  		return nil, nil, true
    89  	}
    90  
    91  	con = p.queue[0]
    92  	// The underlying array will still exist, despite the slice changing, so the object may not GC without this
    93  	// See https://github.com/grpc/grpc-go/issues/4758
    94  	p.queue[0] = nil
    95  	p.queue = p.queue[1:]
    96  
    97  	request = p.pending[con]
    98  	delete(p.pending, con)
    99  
   100  	// Mark the connection as in progress
   101  	p.processing[con] = nil
   102  
   103  	return con, request, false
   104  }
   105  
   106  func (p *PushQueue) MarkDone(con *Connection) {
   107  	p.cond.L.Lock()
   108  	defer p.cond.L.Unlock()
   109  	request := p.processing[con]
   110  	delete(p.processing, con)
   111  
   112  	// If the info is present, that means Enqueue was called while connection was not yet marked done.
   113  	// This means we need to add it back to the queue.
   114  	if request != nil {
   115  		p.pending[con] = request
   116  		p.queue = append(p.queue, con)
   117  		p.cond.Signal()
   118  	}
   119  }
   120  
   121  // Get number of pending proxies
   122  func (p *PushQueue) Pending() int {
   123  	p.cond.L.Lock()
   124  	defer p.cond.L.Unlock()
   125  	return len(p.queue)
   126  }
   127  
   128  // ShutDown will cause queue to ignore all new items added to it. As soon as the
   129  // worker goroutines have drained the existing items in the queue, they will be
   130  // instructed to exit.
   131  func (p *PushQueue) ShutDown() {
   132  	p.cond.L.Lock()
   133  	defer p.cond.L.Unlock()
   134  	p.shuttingDown = true
   135  	p.cond.Broadcast()
   136  }