github.com/uber/kraken@v0.1.4/utils/dedup/limiter.go (about)

     1  // Copyright (c) 2016-2019 Uber Technologies, 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  //     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  package dedup
    15  
    16  import (
    17  	"sync"
    18  	"time"
    19  
    20  	"github.com/andres-erbsen/clock"
    21  )
    22  
    23  // TaskGCInterval is the interval in which garbage collection of old tasks runs.
    24  const TaskGCInterval = time.Minute
    25  
    26  // TaskRunner runs against some input and produces some output w/ a ttl.
    27  type TaskRunner interface {
    28  	Run(input interface{}) (output interface{}, ttl time.Duration)
    29  }
    30  
    31  type task struct {
    32  	input interface{}
    33  
    34  	cond      *sync.Cond
    35  	running   bool
    36  	output    interface{}
    37  	expiresAt time.Time
    38  }
    39  
    40  func newTask(input interface{}) *task {
    41  	return &task{
    42  		input: input,
    43  		cond:  sync.NewCond(new(sync.Mutex)),
    44  	}
    45  }
    46  
    47  func (t *task) expired(now time.Time) bool {
    48  	return now.After(t.expiresAt)
    49  }
    50  
    51  // Limiter deduplicates the running of a common task within a given limit. Tasks
    52  // are deduplicated based on input equality.
    53  type Limiter struct {
    54  	sync.RWMutex
    55  	clk    clock.Clock
    56  	runner TaskRunner
    57  	tasks  map[interface{}]*task
    58  	gc     *IntervalTrap
    59  }
    60  
    61  // NewLimiter creates a new Limiter for tasks. The limit is determined per task
    62  // via the TaskRunner.
    63  func NewLimiter(clk clock.Clock, runner TaskRunner) *Limiter {
    64  	l := &Limiter{
    65  		clk:    clk,
    66  		runner: runner,
    67  		tasks:  make(map[interface{}]*task),
    68  	}
    69  	l.gc = NewIntervalTrap(TaskGCInterval, clk, &limiterTaskGC{l})
    70  	return l
    71  }
    72  
    73  // Run runs a task with input.
    74  func (l *Limiter) Run(input interface{}) interface{} {
    75  	l.gc.Trap()
    76  
    77  	l.RLock()
    78  	t, ok := l.tasks[input]
    79  	l.RUnlock()
    80  	if !ok {
    81  		// Slow path, must initialize task struct under global write lock.
    82  		l.Lock()
    83  		t, ok = l.tasks[input]
    84  		if !ok {
    85  			t = newTask(input)
    86  			l.tasks[input] = t
    87  		}
    88  		l.Unlock()
    89  	}
    90  	return l.getOutput(t)
    91  }
    92  
    93  func (l *Limiter) getOutput(t *task) interface{} {
    94  	t.cond.L.Lock()
    95  
    96  	if !t.expired(l.clk.Now()) {
    97  		defer t.cond.L.Unlock()
    98  		return t.output
    99  	}
   100  
   101  	if t.running {
   102  		t.cond.Wait()
   103  		defer t.cond.L.Unlock()
   104  		return t.output
   105  	}
   106  
   107  	t.running = true
   108  	t.cond.L.Unlock()
   109  
   110  	output, ttl := l.runner.Run(t.input)
   111  
   112  	t.cond.L.Lock()
   113  	t.output = output
   114  	t.expiresAt = l.clk.Now().Add(ttl)
   115  	t.running = false
   116  	t.cond.L.Unlock()
   117  
   118  	t.cond.Broadcast()
   119  
   120  	return output
   121  }
   122  
   123  type limiterTaskGC struct {
   124  	limiter *Limiter
   125  }
   126  
   127  func (gc *limiterTaskGC) Run() {
   128  	gc.limiter.Lock()
   129  	defer gc.limiter.Unlock()
   130  
   131  	for input, t := range gc.limiter.tasks {
   132  		t.cond.L.Lock()
   133  		expired := t.expired(gc.limiter.clk.Now()) && !t.running
   134  		t.cond.L.Unlock()
   135  		if expired {
   136  			delete(gc.limiter.tasks, input)
   137  		}
   138  	}
   139  }