github.com/dubbogo/gost@v1.14.0/time/wheel.go (about)

     1  /*
     2   * Licensed to the Apache Software Foundation (ASF) under one or more
     3   * contributor license agreements.  See the NOTICE file distributed with
     4   * this work for additional information regarding copyright ownership.
     5   * The ASF licenses this file to You under the Apache License, Version 2.0
     6   * (the "License"); you may not use this file except in compliance with
     7   * the License.  You may obtain a copy of the License at
     8   *
     9   *     http://www.apache.org/licenses/LICENSE-2.0
    10   *
    11   * Unless required by applicable law or agreed to in writing, software
    12   * distributed under the License is distributed on an "AS IS" BASIS,
    13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14   * See the License for the specific language governing permissions and
    15   * limitations under the License.
    16   */
    17  
    18  // Package gxtime encapsulates some golang.time functions
    19  // ref: https://github.com/AlexStocks/go-practice/blob/master/time/siddontang_time_wheel.go
    20  package gxtime
    21  
    22  import (
    23  	"sync"
    24  	"time"
    25  )
    26  
    27  type Wheel struct {
    28  	sync.RWMutex
    29  	span   time.Duration
    30  	period time.Duration
    31  	ticker *time.Ticker
    32  	index  int
    33  	ring   []chan struct{}
    34  	once   sync.Once
    35  	now    time.Time
    36  }
    37  
    38  func NewWheel(span time.Duration, buckets int) *Wheel {
    39  	var w *Wheel
    40  
    41  	if span == 0 {
    42  		panic("@span == 0")
    43  	}
    44  	if buckets == 0 {
    45  		panic("@bucket == 0")
    46  	}
    47  
    48  	w = &Wheel{
    49  		span:   span,
    50  		period: span * (time.Duration(buckets)),
    51  		ticker: time.NewTicker(span),
    52  		index:  0,
    53  		ring:   make([]chan struct{}, buckets),
    54  		now:    time.Now(),
    55  	}
    56  
    57  	go func() {
    58  		var notify chan struct{}
    59  		for t := range w.ticker.C {
    60  			w.Lock()
    61  			w.now = t
    62  
    63  			notify = w.ring[w.index]
    64  			w.ring[w.index] = nil
    65  			w.index = (w.index + 1) % len(w.ring)
    66  
    67  			w.Unlock()
    68  
    69  			if notify != nil {
    70  				close(notify)
    71  			}
    72  		}
    73  	}()
    74  
    75  	return w
    76  }
    77  
    78  func (w *Wheel) Stop() {
    79  	w.once.Do(func() { w.ticker.Stop() })
    80  }
    81  
    82  func (w *Wheel) After(timeout time.Duration) <-chan struct{} {
    83  	if timeout >= w.period {
    84  		panic("@timeout over ring's life period")
    85  	}
    86  
    87  	pos := int(timeout / w.span)
    88  	if 0 < pos {
    89  		pos--
    90  	}
    91  
    92  	w.Lock()
    93  	pos = (w.index + pos) % len(w.ring)
    94  	if w.ring[pos] == nil {
    95  		w.ring[pos] = make(chan struct{})
    96  	}
    97  	c := w.ring[pos]
    98  	w.Unlock()
    99  
   100  	return c
   101  }
   102  
   103  func (w *Wheel) Now() time.Time {
   104  	w.RLock()
   105  	now := w.now
   106  	w.RUnlock()
   107  
   108  	return now
   109  }