github.com/15mga/kiwi@v0.0.2-0.20240324021231-b95d5c3ac751/wait.go (about)

     1  package kiwi
     2  
     3  import (
     4  	"github.com/15mga/kiwi/util"
     5  	"os"
     6  	"os/signal"
     7  	"sync"
     8  	"syscall"
     9  	"time"
    10  )
    11  
    12  func WaitGroup(secs int, wg *sync.WaitGroup, ticker util.FnInt) bool {
    13  	tc := time.NewTicker(time.Second)
    14  	ch := make(chan struct{})
    15  	defer func() {
    16  		tc.Stop()
    17  	}()
    18  	go func() {
    19  		wg.Wait()
    20  		close(ch)
    21  	}()
    22  	for {
    23  		select {
    24  		case <-ch:
    25  			return true
    26  		case <-tc.C:
    27  			secs--
    28  			if secs == 0 {
    29  				return false
    30  			}
    31  			if ticker != nil {
    32  				ticker(secs)
    33  			}
    34  		}
    35  	}
    36  }
    37  
    38  func GoWaitGroup(secs int, wg *sync.WaitGroup, over util.FnBool, ticker util.FnInt) {
    39  	go over(WaitGroup(secs, wg, ticker))
    40  }
    41  
    42  type waitInfo struct {
    43  	name string
    44  	fn   util.Fn
    45  }
    46  
    47  var (
    48  	_WaitExitInfos = make([]*waitInfo, 0, 1)
    49  )
    50  
    51  func BeforeExitFn(name string, fn util.Fn) {
    52  	_WaitExitInfos = append(_WaitExitInfos, &waitInfo{
    53  		name: name,
    54  		fn:   fn,
    55  	})
    56  }
    57  
    58  func BeforeExitCh(name string) chan<- struct{} {
    59  	ch := make(chan struct{})
    60  	BeforeExitFn(name, func() {
    61  		<-ch
    62  	})
    63  	return ch
    64  }
    65  
    66  func WaitExit() {
    67  	signalCh := make(chan os.Signal, 1)
    68  	signal.Notify(signalCh, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
    69  	select {
    70  	case <-util.Ctx().Done():
    71  		Info("context done", nil)
    72  	case s := <-signalCh:
    73  		Info("signal notify", util.M{
    74  			"signal": s,
    75  		})
    76  		util.Cancel()
    77  	}
    78  
    79  	waitCh := make(chan struct{})
    80  	go func() {
    81  		count := len(_WaitExitInfos)
    82  		if count == 0 {
    83  			close(waitCh)
    84  			return
    85  		}
    86  		nameCh := make(chan string, count)
    87  		status := make(util.M, count)
    88  		for _, info := range _WaitExitInfos {
    89  			status[info.name] = false
    90  			wInfo := info
    91  			go func() {
    92  				Info("wait exit", util.M{
    93  					"name": wInfo.name,
    94  				})
    95  				wInfo.fn()
    96  				nameCh <- wInfo.name
    97  			}()
    98  		}
    99  		ticker := time.NewTicker(time.Second)
   100  		tickCount := 0
   101  		for {
   102  			select {
   103  			case <-ticker.C:
   104  				Info("exit status", util.M{
   105  					"status": status,
   106  					"count":  tickCount,
   107  				})
   108  				tickCount++
   109  			case name := <-nameCh:
   110  				Info("exit", util.M{
   111  					"name": name,
   112  				})
   113  				status[name] = true
   114  				count--
   115  				if count == 0 {
   116  					ticker.Stop()
   117  					close(waitCh)
   118  					return
   119  				}
   120  			}
   121  		}
   122  	}()
   123  
   124  	timeout := time.NewTimer(time.Second * 60)
   125  	select {
   126  	case <-timeout.C:
   127  		Info("exit timeout", nil)
   128  	case <-waitCh:
   129  		timeout.Stop()
   130  		Info("exit complete", nil)
   131  	}
   132  }