github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/transport/collect.go (about)

     1  // Package transport provides long-lived http/tcp connections for
     2  // intra-cluster communications (see README for details and usage example).
     3  /*
     4   * Copyright (c) 2018-2024, NVIDIA CORPORATION. All rights reserved.
     5   */
     6  package transport
     7  
     8  import (
     9  	"container/heap"
    10  	"time"
    11  
    12  	"github.com/NVIDIA/aistore/cmn/atomic"
    13  	"github.com/NVIDIA/aistore/cmn/cos"
    14  	"github.com/NVIDIA/aistore/cmn/debug"
    15  	"github.com/NVIDIA/aistore/cmn/nlog"
    16  )
    17  
    18  type (
    19  	ctrl struct { // add/del channel to/from collector
    20  		s   *streamBase
    21  		add bool
    22  	}
    23  	collector struct {
    24  		streams map[string]*streamBase
    25  		ticker  *time.Ticker
    26  		stopCh  cos.StopCh
    27  		ctrlCh  chan ctrl
    28  		heap    []*streamBase
    29  		none    atomic.Bool // no streams
    30  	}
    31  )
    32  
    33  var (
    34  	sc *StreamCollector // idle timer and house-keeping (slow path)
    35  	gc *collector       // real stream collector
    36  )
    37  
    38  // interface guard
    39  var _ cos.Runner = (*StreamCollector)(nil)
    40  
    41  // Stream Collector:
    42  // 1. controls stream activation (followed by connection establishment and HTTP PUT), and
    43  //    deactivation (teardown)
    44  // 2. provides each stream with its own idle timer (with timeout measured in ticks - see tickUnit)
    45  // 3. deactivates idle streams
    46  
    47  func (*StreamCollector) Name() string { return "stream-collector" }
    48  
    49  func (sc *StreamCollector) Run() error {
    50  	cos.Infof("Intra-cluster networking: %s client", whichClient())
    51  	cos.Infof("Starting %s", sc.Name())
    52  	gc.ticker = time.NewTicker(dfltTickIdle)
    53  	gc.none.Store(true)
    54  	gc.run()
    55  	gc.ticker.Stop()
    56  	return nil
    57  }
    58  
    59  func (sc *StreamCollector) Stop(err error) {
    60  	nlog.Infoln("Stopping", sc.Name(), "err:", err)
    61  	gc.stop()
    62  }
    63  
    64  func (gc *collector) run() {
    65  	for {
    66  		select {
    67  		case <-gc.ticker.C:
    68  			gc.do()
    69  		case ctrl, ok := <-gc.ctrlCh:
    70  			if !ok {
    71  				return
    72  			}
    73  			s, add := ctrl.s, ctrl.add
    74  			_, ok = gc.streams[s.lid]
    75  			if add {
    76  				debug.Assert(!ok, s.lid)
    77  				gc.streams[s.lid] = s
    78  				heap.Push(gc, s)
    79  				if gc.none.CAS(true, false) {
    80  					gc.ticker.Reset(dfltTick)
    81  				}
    82  			} else if ok {
    83  				heap.Remove(gc, s.time.index)
    84  				s.time.ticks = 1
    85  			}
    86  		case <-gc.stopCh.Listen():
    87  			for _, s := range gc.streams {
    88  				s.Stop()
    89  			}
    90  			gc.streams = nil
    91  			return
    92  		}
    93  	}
    94  }
    95  
    96  func (gc *collector) stop() {
    97  	gc.stopCh.Close()
    98  }
    99  
   100  func (gc *collector) remove(s *streamBase) {
   101  	gc.ctrlCh <- ctrl{s, false} // remove and close workCh
   102  }
   103  
   104  // as min-heap
   105  func (gc *collector) Len() int { return len(gc.heap) }
   106  
   107  func (gc *collector) Less(i, j int) bool {
   108  	si := gc.heap[i]
   109  	sj := gc.heap[j]
   110  	return si.time.ticks < sj.time.ticks
   111  }
   112  
   113  func (gc *collector) Swap(i, j int) {
   114  	gc.heap[i], gc.heap[j] = gc.heap[j], gc.heap[i]
   115  	gc.heap[i].time.index = i
   116  	gc.heap[j].time.index = j
   117  }
   118  
   119  func (gc *collector) Push(x any) {
   120  	l := len(gc.heap)
   121  	s := x.(*streamBase)
   122  	s.time.index = l
   123  	gc.heap = append(gc.heap, s)
   124  	heap.Fix(gc, s.time.index) // reorder the newly added stream right away
   125  }
   126  
   127  func (gc *collector) update(s *streamBase, ticks int) {
   128  	s.time.ticks = ticks
   129  	debug.Assert(s.time.ticks >= 0)
   130  	heap.Fix(gc, s.time.index)
   131  }
   132  
   133  func (gc *collector) Pop() any {
   134  	old := gc.heap
   135  	n := len(old)
   136  	sl := old[n-1]
   137  	gc.heap = old[0 : n-1]
   138  	return sl
   139  }
   140  
   141  // collector's main method
   142  func (gc *collector) do() {
   143  	for lid, s := range gc.streams {
   144  		if s.IsTerminated() {
   145  			_, err := s.TermInfo()
   146  			if s.time.inSend.Swap(false) {
   147  				s.streamer.drain(err)
   148  				s.time.ticks = 1
   149  				continue
   150  			}
   151  
   152  			s.time.ticks--
   153  			if s.time.ticks <= 0 {
   154  				delete(gc.streams, lid)
   155  				if len(gc.streams) == 0 {
   156  					gc.ticker.Reset(dfltTickIdle)
   157  					debug.Assert(!gc.none.Load())
   158  					gc.none.Store(true)
   159  				}
   160  				s.streamer.closeAndFree()
   161  				s.streamer.abortPending(err, true /*completions*/)
   162  			}
   163  		} else if s.sessST.Load() == active {
   164  			gc.update(s, s.time.ticks-1)
   165  		}
   166  	}
   167  	for _, s := range gc.streams {
   168  		if s.time.ticks > 0 {
   169  			continue
   170  		}
   171  		gc.update(s, int(s.time.idleTeardown/dfltTick))
   172  		if s.time.inSend.Swap(false) {
   173  			continue
   174  		}
   175  		s.streamer.idleTick()
   176  	}
   177  	// at this point the following must be true for each i = range gc.heap:
   178  	// 1. heap[i].index == i
   179  	// 2. heap[i+1].ticks >= heap[i].ticks
   180  }