github.com/TeaOSLab/EdgeNode@v1.3.8/internal/utils/ratelimit/bandwidth.go (about)

     1  // Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
     2  
     3  package ratelimit
     4  
     5  import (
     6  	"context"
     7  	"github.com/TeaOSLab/EdgeNode/internal/utils/fasttime"
     8  	"sync/atomic"
     9  	"time"
    10  )
    11  
    12  // Bandwidth lossy bandwidth limiter
    13  type Bandwidth struct {
    14  	totalBytes int64
    15  
    16  	currentTimestamp int64
    17  	currentBytes     int64
    18  }
    19  
    20  // NewBandwidth create new bandwidth limiter
    21  func NewBandwidth(totalBytes int64) *Bandwidth {
    22  	return &Bandwidth{totalBytes: totalBytes}
    23  }
    24  
    25  // Ack acquire next chance to send data
    26  func (this *Bandwidth) Ack(ctx context.Context, newBytes int) {
    27  	if newBytes <= 0 {
    28  		return
    29  	}
    30  	if this.totalBytes <= 0 {
    31  		return
    32  	}
    33  
    34  	var timestamp = fasttime.Now().Unix()
    35  	if this.currentTimestamp != 0 && this.currentTimestamp != timestamp {
    36  		this.currentTimestamp = timestamp
    37  		this.currentBytes = int64(newBytes)
    38  
    39  		// 第一次发送直接放行，不需要判断
    40  		return
    41  	}
    42  
    43  	if this.currentTimestamp == 0 {
    44  		this.currentTimestamp = timestamp
    45  	}
    46  	if atomic.AddInt64(&this.currentBytes, int64(newBytes)) <= this.totalBytes {
    47  		return
    48  	}
    49  
    50  	var timeout = time.NewTimer(1 * time.Second)
    51  	if ctx != nil {
    52  		select {
    53  		case <-timeout.C:
    54  		case <-ctx.Done():
    55  		}
    56  	} else {
    57  		select {
    58  		case <-timeout.C:
    59  		}
    60  	}
    61  }