github.com/livekit/protocol@v1.16.1-0.20240517185851-47e4c6bba773/utils/protoproxy.go (about)

     1  /*
     2   * Copyright 2023 LiveKit, Inc
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * 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  
    17  package utils
    18  
    19  import (
    20  	"sync"
    21  	"time"
    22  
    23  	"github.com/frostbyte73/core"
    24  	"google.golang.org/protobuf/proto"
    25  )
    26  
    27  // ProtoProxy is a caching proxy for protobuf messages that may be expensive
    28  // to compute. It is used to avoid unnecessary re-generation of Protobufs
    29  type ProtoProxy[T proto.Message] struct {
    30  	message         T
    31  	updateFn        func() T
    32  	fuse            core.Fuse
    33  	updateChan      chan struct{}
    34  	awaitChan       chan struct{}
    35  	done            chan struct{}
    36  	queueUpdate     chan struct{}
    37  	dirty           bool
    38  	refreshedAt     time.Time
    39  	refreshInterval time.Duration
    40  	lock            sync.RWMutex
    41  }
    42  
    43  // NewProtoProxy creates a new ProtoProxy that regenerates underlying values at a cadence of refreshInterval
    44  // this should be used for updates that should be sent periodically, but does not have the urgency of immediate delivery
    45  // updateFn should provide computations required to generate the protobuf
    46  // if refreshInterval is 0, then proxy will only update on MarkDirty(true)
    47  func NewProtoProxy[T proto.Message](refreshInterval time.Duration, updateFn func() T) *ProtoProxy[T] {
    48  	p := &ProtoProxy[T]{
    49  		updateChan:      make(chan struct{}, 1),
    50  		updateFn:        updateFn,
    51  		done:            make(chan struct{}),
    52  		refreshInterval: refreshInterval,
    53  		queueUpdate:     make(chan struct{}, 1),
    54  	}
    55  	p.performUpdate(true)
    56  	if refreshInterval > 0 {
    57  		go p.worker()
    58  	}
    59  	return p
    60  }
    61  
    62  func (p *ProtoProxy[T]) MarkDirty(immediate bool) <-chan struct{} {
    63  	p.lock.Lock()
    64  	p.dirty = true
    65  	shouldUpdate := immediate || time.Since(p.refreshedAt) > p.refreshInterval
    66  
    67  	if p.awaitChan == nil {
    68  		p.awaitChan = make(chan struct{})
    69  	}
    70  	awaitChan := p.awaitChan
    71  	p.lock.Unlock()
    72  
    73  	if shouldUpdate {
    74  		select {
    75  		case p.queueUpdate <- struct{}{}:
    76  		default:
    77  		}
    78  	}
    79  	return awaitChan
    80  }
    81  
    82  func (p *ProtoProxy[T]) Updated() <-chan struct{} {
    83  	return p.updateChan
    84  }
    85  
    86  func (p *ProtoProxy[T]) Get() T {
    87  	p.lock.RLock()
    88  	defer p.lock.RUnlock()
    89  	return proto.Clone(p.message).(T)
    90  }
    91  
    92  func (p *ProtoProxy[T]) Stop() {
    93  	p.fuse.Break()
    94  
    95  	// goroutine is not started when refreshInterval == 0
    96  	if p.refreshInterval > 0 {
    97  		<-p.done
    98  	}
    99  
   100  	p.lock.Lock()
   101  	defer p.lock.Unlock()
   102  	if awaitChan := p.awaitChan; awaitChan != nil {
   103  		p.awaitChan = nil
   104  		close(awaitChan)
   105  	}
   106  }
   107  
   108  func (p *ProtoProxy[T]) performUpdate(skipNotify bool) {
   109  	// set dirty back *before* calling updateFn because otherwise it could
   110  	// wipe out another thread setting dirty to true while updateFn is executing
   111  	p.lock.Lock()
   112  	p.dirty = false
   113  
   114  	if awaitChan := p.awaitChan; awaitChan != nil {
   115  		p.awaitChan = nil
   116  		defer close(awaitChan)
   117  	}
   118  	p.lock.Unlock()
   119  
   120  	msg := p.updateFn()
   121  
   122  	p.lock.Lock()
   123  	if proto.Equal(p.message, msg) {
   124  		// no change, skip the notification
   125  		p.lock.Unlock()
   126  		return
   127  	}
   128  	p.message = msg
   129  	// only updating refreshedAt if we have notified, so it shouldn't push
   130  	// out the next notification out by another interval
   131  	p.refreshedAt = time.Now()
   132  	p.lock.Unlock()
   133  
   134  	if !skipNotify {
   135  		select {
   136  		case p.updateChan <- struct{}{}:
   137  		default:
   138  		}
   139  	}
   140  }
   141  
   142  func (p *ProtoProxy[T]) worker() {
   143  	ticker := time.NewTicker(p.refreshInterval)
   144  	defer ticker.Stop()
   145  	defer close(p.done)
   146  
   147  	for {
   148  		select {
   149  		case <-p.fuse.Watch():
   150  			return
   151  		case <-ticker.C:
   152  			p.lock.RLock()
   153  			shouldUpdate := p.dirty && time.Since(p.refreshedAt) > p.refreshInterval
   154  			p.lock.RUnlock()
   155  			if shouldUpdate {
   156  				p.performUpdate(false)
   157  			}
   158  		case <-p.queueUpdate:
   159  			p.performUpdate(false)
   160  		}
   161  	}
   162  }