github.com/machinefi/w3bstream@v1.6.5-rc9.0.20240426031326-b8c7c4876e72/pkg/modules/trafficlimit/traffic_scheduler.go (about)

     1  package trafficlimit
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"time"
     7  
     8  	"github.com/go-co-op/gocron/v2"
     9  	"github.com/pkg/errors"
    10  
    11  	"github.com/machinefi/w3bstream/pkg/depends/conf/logger"
    12  	"github.com/machinefi/w3bstream/pkg/depends/conf/redis"
    13  	"github.com/machinefi/w3bstream/pkg/depends/x/mapx"
    14  	"github.com/machinefi/w3bstream/pkg/models"
    15  	"github.com/machinefi/w3bstream/pkg/types"
    16  )
    17  
    18  func AddScheduler(ctx context.Context, v *models.TrafficLimit, start ...bool) error {
    19  	_, l := logger.NewSpanContext(ctx, "traffic_limit.AddScheduler")
    20  	defer l.End()
    21  
    22  	l = l.WithValues("traffic_limit", v.TrafficLimitID, "du", v.Duration, "count", v.Threshold)
    23  	s, ok := schedulers.Load(v.TrafficLimitID)
    24  	if !ok {
    25  		s, err := NewScheduler(ctx, v, start...)
    26  		if err != nil {
    27  			l.Error(err)
    28  			return err
    29  		}
    30  		schedulers.Store(v.TrafficLimitID, s)
    31  		l.Info("schedule created")
    32  		return nil
    33  	}
    34  	s.update(ctx, v)
    35  	return nil
    36  }
    37  
    38  func AddAndStartScheduler(ctx context.Context, v *models.TrafficLimit) error {
    39  	return AddScheduler(ctx, v, true)
    40  }
    41  
    42  func RmvScheduler(ctx context.Context, id types.SFID) {
    43  	_, l := logger.NewSpanContext(ctx, "traffic_limit.RmvScheduler")
    44  	defer l.End()
    45  
    46  	s, _ := schedulers.Load(id)
    47  	if s != nil {
    48  		s.Stop()
    49  	}
    50  	schedulers.Remove(id)
    51  	l.WithValues("traffic_limit", id).Info("schedule removed")
    52  }
    53  
    54  var schedulers = *mapx.New[types.SFID, *Scheduler]()
    55  
    56  func NewScheduler(ctx context.Context, v *models.TrafficLimit, start ...bool) (*Scheduler, error) {
    57  	s := &Scheduler{
    58  		key: v.CacheKey(),
    59  		du:  v.Duration.Duration(),
    60  		cnt: int64(v.Threshold),
    61  		kv:  types.MustRedisEndpointFromContext(ctx).WithPrefix(prefix),
    62  	}
    63  	var err error
    64  	s.sch, err = gocron.NewScheduler()
    65  	if err != nil {
    66  		return nil, err
    67  	}
    68  	_, err = s.sch.NewJob(gocron.DurationJob(s.du), gocron.NewTask(s.reset))
    69  	if err != nil {
    70  		return nil, err
    71  	}
    72  	if len(start) > 0 && start[0] {
    73  		s.Start()
    74  	}
    75  	return s, nil
    76  }
    77  
    78  type Scheduler struct {
    79  	mu  sync.Mutex
    80  	key string
    81  	du  time.Duration
    82  	cnt int64
    83  	sch gocron.Scheduler
    84  	kv  *redis.Redis
    85  }
    86  
    87  func (s *Scheduler) Start() {
    88  	s.reset()
    89  	s.sch.Start()
    90  }
    91  
    92  func (s *Scheduler) Stop() {
    93  	_ = s.kv.Del(s.key)
    94  	_ = s.sch.Shutdown()
    95  }
    96  
    97  func (s *Scheduler) reset() {
    98  	s.mu.Lock()
    99  	key, cnt, du := s.key, s.cnt, s.du
   100  	s.mu.Unlock()
   101  	_ = s.kv.SetEx(key, cnt, du)
   102  }
   103  
   104  func (s *Scheduler) update(ctx context.Context, v *models.TrafficLimit) {
   105  	_, l := logger.NewSpanContext(ctx, "traffic_limit.AddScheduler")
   106  	defer l.End()
   107  
   108  	l = l.WithValues("traffic_limit", v.TrafficLimitID)
   109  
   110  	s.mu.Lock()
   111  	defer s.mu.Unlock()
   112  
   113  	if s.key != v.CacheKey() || s.cnt != int64(v.Threshold) || s.du != v.Duration.Duration() {
   114  		s.key = v.CacheKey()
   115  		s.du = v.Duration.Duration()
   116  		s.cnt = int64(v.Threshold)
   117  		err := s.kv.SetEx(s.key, s.cnt, s.du)
   118  		l = l.WithValues("traffic_limit", v.TrafficLimitInfo, "du", v.Duration, "count", s.cnt)
   119  		if err != nil {
   120  			l.Error(errors.Wrap(err, "schedule update failed: %v"))
   121  		} else {
   122  			l.Info("schedule updated")
   123  		}
   124  	}
   125  }