github.com/livekit/protocol@v1.5.7/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  	done            chan struct{}
    35  	queueUpdate     chan struct{}
    36  	dirty           bool
    37  	refreshedAt     time.Time
    38  	refreshInterval time.Duration
    39  	lock            sync.RWMutex
    40  }
    41  
    42  // NewProtoProxy creates a new ProtoProxy that regenerates underlying values at a cadence of refreshInterval
    43  // this should be used for updates that should be sent periodically, but does not have the urgency of immediate delivery
    44  // updateFn should provide computations required to generate the protobuf
    45  // if refreshInterval is 0, then proxy will only update on MarkDirty(true)
    46  func NewProtoProxy[T proto.Message](refreshInterval time.Duration, updateFn func() T) *ProtoProxy[T] {
    47  	p := &ProtoProxy[T]{
    48  		updateChan:      make(chan struct{}, 1),
    49  		updateFn:        updateFn,
    50  		done:            make(chan struct{}),
    51  		fuse:            core.NewFuse(),
    52  		refreshInterval: refreshInterval,
    53  	}
    54  	p.performUpdate(true)
    55  	if refreshInterval > 0 {
    56  		go p.worker()
    57  	}
    58  	return p
    59  }
    60  
    61  func (p *ProtoProxy[T]) MarkDirty(immediate bool) {
    62  	p.lock.Lock()
    63  	p.dirty = true
    64  	shouldUpdate := immediate || time.Since(p.refreshedAt) > p.refreshInterval
    65  	p.lock.Unlock()
    66  	if shouldUpdate {
    67  		select {
    68  		case p.queueUpdate <- struct{}{}:
    69  		default:
    70  		}
    71  	}
    72  }
    73  
    74  func (p *ProtoProxy[T]) Updated() <-chan struct{} {
    75  	return p.updateChan
    76  }
    77  
    78  func (p *ProtoProxy[T]) Get() T {
    79  	p.lock.RLock()
    80  	defer p.lock.RUnlock()
    81  	return p.message
    82  }
    83  
    84  func (p *ProtoProxy[T]) Stop() {
    85  	p.fuse.Break()
    86  
    87  	// goroutine is not started when refreshInterval == 0
    88  	if p.refreshInterval > 0 {
    89  		<-p.done
    90  	}
    91  }
    92  
    93  func (p *ProtoProxy[T]) performUpdate(skipNotify bool) {
    94  	msg := p.updateFn()
    95  	p.lock.Lock()
    96  	p.message = msg
    97  	p.refreshedAt = time.Now()
    98  	p.dirty = false
    99  	p.lock.Unlock()
   100  
   101  	if !skipNotify {
   102  		select {
   103  		case p.updateChan <- struct{}{}:
   104  		default:
   105  		}
   106  	}
   107  }
   108  
   109  func (p *ProtoProxy[T]) worker() {
   110  	ticker := time.NewTicker(p.refreshInterval)
   111  	defer ticker.Stop()
   112  	defer close(p.done)
   113  
   114  	for {
   115  		select {
   116  		case <-p.fuse.Watch():
   117  			return
   118  		case <-ticker.C:
   119  			p.lock.RLock()
   120  			shouldUpdate := p.dirty && time.Since(p.refreshedAt) > p.refreshInterval
   121  			p.lock.RUnlock()
   122  			if shouldUpdate {
   123  				p.performUpdate(false)
   124  			}
   125  		case <-p.queueUpdate:
   126  			p.performUpdate(false)
   127  		}
   128  	}
   129  }