github.com/Mrs4s/MiraiGo@v0.0.0-20240226124653-54bdd873e3fe/utils/waiter.go (about)

     1  package utils
     2  
     3  import "sync"
     4  
     5  // UploadWaiter 用于控制并发上传，当有一个文件多次上传时，
     6  // 等待第一个上传，后续的上传并发进行(可以秒传).
     7  type UploadWaiter struct {
     8  	mu sync.Mutex
     9  	m  map[string]*sync.WaitGroup
    10  }
    11  
    12  // NewUploadWaiter return a new UploadWaiter.
    13  func NewUploadWaiter() *UploadWaiter {
    14  	return &UploadWaiter{
    15  		m: make(map[string]*sync.WaitGroup),
    16  	}
    17  }
    18  
    19  // Wait 如果不是第一个上传则等待。
    20  func (s *UploadWaiter) Wait(key string) {
    21  	s.mu.Lock()
    22  	if w, ok := s.m[key]; ok {
    23  		s.mu.Unlock()
    24  		w.Wait()
    25  	} else {
    26  		wg := new(sync.WaitGroup)
    27  		wg.Add(1)
    28  		s.m[key] = wg
    29  		s.mu.Unlock()
    30  	}
    31  }
    32  
    33  // Done 当前上传任务已完成。
    34  func (s *UploadWaiter) Done(key string) {
    35  	s.mu.Lock()
    36  	if w, ok := s.m[key]; ok {
    37  		w.Done()
    38  		delete(s.m, key)
    39  	}
    40  	s.mu.Unlock()
    41  }