server_actors.go raw

   1  package app
   2  
   3  import (
   4  	"time"
   5  
   6  	"git.smesh.lol/actor"
   7  )
   8  
   9  // -- connPerIP actor types --
  10  
  11  type connIPCheckArgs struct {
  12  	IP  string
  13  	Max int
  14  }
  15  type connIPCheckResp struct {
  16  	allowed bool
  17  	current int
  18  }
  19  
  20  func (s *Server) startConnPerIPActor() {
  21  	s.connIPCheck = actor.NewFunc[connIPCheckArgs, connIPCheckResp]()
  22  	s.connIPDec = actor.NewInbox[string](16)
  23  	s.connIPLC = actor.NewLifecycle()
  24  	actor.Go(s.connIPLC, func() {
  25  		m := make(map[string]int)
  26  		for {
  27  			select {
  28  			case msg := <-s.connIPCheck.Recv():
  29  				current := m[msg.Req.IP]
  30  				if current >= msg.Req.Max {
  31  					msg.Reply(connIPCheckResp{allowed: false, current: current})
  32  				} else {
  33  					m[msg.Req.IP]++
  34  					msg.Reply(connIPCheckResp{allowed: true, current: current + 1})
  35  				}
  36  			case ip := <-s.connIPDec.Recv():
  37  				m[ip]--
  38  				if m[ip] <= 0 {
  39  					delete(m, ip)
  40  				}
  41  			case <-s.Ctx.Done():
  42  				return
  43  			}
  44  		}
  45  	})
  46  }
  47  
  48  func (s *Server) ConnIPCheckAndInc(ip string, max int) (allowed bool, current int) {
  49  	r := s.connIPCheck.Call(connIPCheckArgs{IP: ip, Max: max})
  50  	return r.allowed, r.current
  51  }
  52  
  53  func (s *Server) ConnIPDec(ip string) {
  54  	s.connIPDec.TrySend(ip)
  55  }
  56  
  57  // -- challenge actor types --
  58  
  59  type chalSetArgs struct {
  60  	Key  string
  61  	Data []byte
  62  }
  63  type chalGetResp struct {
  64  	data   []byte
  65  	exists bool
  66  }
  67  
  68  func (s *Server) startChallengeActor() {
  69  	s.chalSet = actor.NewInbox[chalSetArgs](4)
  70  	s.chalGet = actor.NewFunc[string, chalGetResp]()
  71  	s.chalDelete = actor.NewInbox[string](4)
  72  	s.chalLC = actor.NewLifecycle()
  73  	actor.Go(s.chalLC, func() {
  74  		m := make(map[string][]byte)
  75  		for {
  76  			select {
  77  			case args := <-s.chalSet.Recv():
  78  				m[args.Key] = args.Data
  79  			case msg := <-s.chalGet.Recv():
  80  				data, exists := m[msg.Req]
  81  				msg.Reply(chalGetResp{data: data, exists: exists})
  82  			case key := <-s.chalDelete.Recv():
  83  				delete(m, key)
  84  			case <-s.Ctx.Done():
  85  				return
  86  			}
  87  		}
  88  	})
  89  }
  90  
  91  func (s *Server) ChallengeSet(key string, data []byte) {
  92  	s.chalSet.TrySend(chalSetArgs{Key: key, Data: data})
  93  }
  94  
  95  func (s *Server) ChallengeGet(key string) ([]byte, bool) {
  96  	r := s.chalGet.Call(key)
  97  	return r.data, r.exists
  98  }
  99  
 100  func (s *Server) ChallengeDelete(key string) {
 101  	s.chalDelete.TrySend(key)
 102  }
 103  
 104  func (s *Server) ChallengeSetWithExpiry(key string, data []byte, ttl time.Duration) {
 105  	s.ChallengeSet(key, data)
 106  	go func() {
 107  		time.Sleep(ttl)
 108  		s.ChallengeDelete(key)
 109  	}()
 110  }
 111  
 112  // -- message gate (drain-then-pause) --
 113  //
 114  // The gate tracks in-flight message handlers via a counting channel.
 115  // enterCh is buffered to the max concurrent handler count.
 116  // PauseMessageProcessing blocks new entries and drains existing ones.
 117  
 118  type msgGateEnterReq struct{}
 119  type msgGateExitReq struct{}
 120  type msgGatePauseReq struct {
 121  	resp chan struct{} // buffered 1: ack when drain complete
 122  }
 123  type msgGateResumeReq struct{}
 124  
 125  func (s *Server) startMessageGate() {
 126  	s.msgGateEnterCh = make(chan msgGateEnterReq, 128) // buffered: absorb bursts
 127  	s.msgGateExitCh = make(chan msgGateExitReq, 128)
 128  	s.msgGatePauseCh = make(chan msgGatePauseReq)
 129  	s.msgGateResumeCh = make(chan msgGateResumeReq)
 130  	s.msgGateDone = make(chan struct{})
 131  	go func() {
 132  		defer close(s.msgGateDone)
 133  		inFlight := 0
 134  		paused := false
 135  		var pauseAck chan struct{} // set when paused and waiting for drain
 136  
 137  		for {
 138  			if paused {
 139  				// While paused, only accept exits and resume
 140  				select {
 141  				case <-s.msgGateExitCh:
 142  					inFlight--
 143  					if inFlight == 0 && pauseAck != nil {
 144  						close(pauseAck)
 145  						pauseAck = nil
 146  					}
 147  				case <-s.msgGateResumeCh:
 148  					paused = false
 149  				case <-s.Ctx.Done():
 150  					return
 151  				}
 152  			} else {
 153  				select {
 154  				case <-s.msgGateEnterCh:
 155  					inFlight++
 156  				case <-s.msgGateExitCh:
 157  					inFlight--
 158  				case req := <-s.msgGatePauseCh:
 159  					paused = true
 160  					if inFlight == 0 {
 161  						close(req.resp)
 162  					} else {
 163  						pauseAck = req.resp
 164  					}
 165  				case <-s.Ctx.Done():
 166  					return
 167  				}
 168  			}
 169  		}
 170  	}()
 171  }
 172  
 173  func (s *Server) PauseMessageProcessing() {
 174  	resp := make(chan struct{}, 1)
 175  	s.msgGatePauseCh <- msgGatePauseReq{resp: resp}
 176  	<-resp // blocks until all in-flight drained
 177  }
 178  
 179  func (s *Server) ResumeMessageProcessing() {
 180  	s.msgGateResumeCh <- msgGateResumeReq{}
 181  }
 182  
 183  func (s *Server) AcquireMessageProcessingLock() {
 184  	s.msgGateEnterCh <- msgGateEnterReq{}
 185  }
 186  
 187  func (s *Server) ReleaseMessageProcessingLock() {
 188  	s.msgGateExitCh <- msgGateExitReq{}
 189  }
 190