github.com/yandex/pandora@v0.5.32/core/schedule/composite.go (about)

     1  package schedule
     2  
     3  import (
     4  	"sync"
     5  	"time"
     6  
     7  	"github.com/yandex/pandora/core"
     8  )
     9  
    10  type CompositeConf struct {
    11  	Nested []core.Schedule `config:"nested"`
    12  }
    13  
    14  func NewCompositeConf(conf CompositeConf) core.Schedule {
    15  	return NewComposite(conf.Nested...)
    16  }
    17  
    18  func NewComposite(scheds ...core.Schedule) core.Schedule {
    19  	switch len(scheds) {
    20  	case 0:
    21  		return NewOnce(0)
    22  	case 1:
    23  		return scheds[0]
    24  	}
    25  
    26  	var (
    27  		left            = make([]int, len(scheds))
    28  		unknown         bool // If meet any Left() < 0, all previous leftBefore is unknown
    29  		leftAccumulator int  // If unknown, then at least leftBefore accumulated, else exactly leftBefore.
    30  	)
    31  	for i := len(scheds) - 1; i >= 0; i-- {
    32  		left[i] = leftAccumulator
    33  		schedLeft := scheds[i].Left()
    34  		if schedLeft < 0 {
    35  			schedLeft = -1
    36  			unknown = true
    37  			leftAccumulator = -1
    38  		}
    39  		if !unknown {
    40  			leftAccumulator += schedLeft
    41  		}
    42  	}
    43  
    44  	return &compositeSchedule{
    45  		scheds:    scheds,
    46  		leftAfter: left,
    47  	}
    48  }
    49  
    50  type compositeSchedule struct {
    51  	// Under read lock, goroutine can read slices, it's values, and call values goroutine safe methods.
    52  	// Under write lock, goroutine can do anything.
    53  	rwMu      sync.RWMutex
    54  	scheds    []core.Schedule // At least once schedule. First schedule can be finished.
    55  	leftAfter []int           // Tokens leftBefore, if known exactly, or at least tokens leftBefore otherwise.
    56  }
    57  
    58  func (s *compositeSchedule) Start(startAt time.Time) {
    59  	s.rwMu.Lock()
    60  	defer s.rwMu.Unlock()
    61  	s.scheds[0].Start(startAt)
    62  }
    63  func (s *compositeSchedule) Next() (tx time.Time, ok bool) {
    64  	s.rwMu.RLock()
    65  	tx, ok = s.scheds[0].Next()
    66  	if ok {
    67  		s.rwMu.RUnlock()
    68  		return // Got token, all is good.
    69  	}
    70  	schedsLeft := len(s.scheds)
    71  	s.rwMu.RUnlock()
    72  	if schedsLeft == 1 {
    73  		return // All nested schedules has been finished, so composite is finished too.
    74  	}
    75  	// Current schedule is finished, but some are left.
    76  	// Let's start next, with got finish time from previous!
    77  	s.rwMu.Lock()
    78  	schedsLeftNow := len(s.scheds)
    79  	somebodyStartedNextBeforeUs := schedsLeftNow < schedsLeft
    80  	if somebodyStartedNextBeforeUs {
    81  		// Let's just take token.
    82  		tx, ok = s.scheds[0].Next()
    83  		s.rwMu.Unlock()
    84  		if ok || schedsLeftNow == 1 {
    85  			return
    86  		}
    87  		// Very strange. Schedule was started and drained while we was waiting for it.
    88  		// Should very rare, so let's just retry.
    89  		return s.Next()
    90  	}
    91  	s.startNext(tx)
    92  	tx, ok = s.scheds[0].Next()
    93  	s.rwMu.Unlock()
    94  	if !ok && schedsLeftNow > 1 {
    95  		// What? Schedule without any tokens? Okay, just retry.
    96  		return s.Next()
    97  	}
    98  	return
    99  }
   100  
   101  func (s *compositeSchedule) Left() int {
   102  	s.rwMu.RLock()
   103  	schedsLeft := len(s.scheds)
   104  	leftAfter := int(s.leftAfter[0])
   105  	left := s.scheds[0].Left()
   106  	s.rwMu.RUnlock()
   107  	if schedsLeft == 1 {
   108  		return left
   109  	}
   110  	if left == 0 {
   111  		if leftAfter >= 0 {
   112  			return leftAfter
   113  		}
   114  		// leftAfter was unknown, at schedule create moment.
   115  		// But now, it can be finished. Let's shift, and try one more time.
   116  		s.rwMu.Lock()
   117  		shedsLeftNow := len(s.scheds)
   118  		if shedsLeftNow == schedsLeft {
   119  			currentFinishTime, ok := s.scheds[0].Next()
   120  			if ok {
   121  				s.rwMu.Unlock()
   122  				panic("current schedule is not finished")
   123  			}
   124  			s.startNext(currentFinishTime)
   125  		}
   126  		s.rwMu.Unlock()
   127  		return s.Left()
   128  	}
   129  	if left < 0 {
   130  		return -1
   131  	}
   132  	return left + leftAfter
   133  }
   134  
   135  func (s *compositeSchedule) startNext(currentFinishTime time.Time) {
   136  	s.scheds = s.scheds[1:]
   137  	s.leftAfter = s.leftAfter[1:]
   138  	s.scheds[0].Start(currentFinishTime)
   139  }