github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/admin/commands/collection/tx_rate_limiter.go (about)

     1  package collection
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  
     7  	"github.com/rs/zerolog/log"
     8  	"golang.org/x/time/rate"
     9  
    10  	"github.com/onflow/flow-go/admin"
    11  	"github.com/onflow/flow-go/admin/commands"
    12  	"github.com/onflow/flow-go/engine/collection/ingest"
    13  )
    14  
    15  var _ commands.AdminCommand = (*TxRateLimitCommand)(nil)
    16  
    17  // TxRateLimitCommand will adjust the transaction ingest rate limiter.
    18  type TxRateLimitCommand struct {
    19  	limiter *ingest.AddressRateLimiter
    20  }
    21  
    22  type TxRateLimitCommandAddress struct {
    23  	Addresses []string
    24  }
    25  
    26  func NewTxRateLimitCommand(limiter *ingest.AddressRateLimiter) *TxRateLimitCommand {
    27  	return &TxRateLimitCommand{
    28  		limiter: limiter,
    29  	}
    30  }
    31  
    32  func (s *TxRateLimitCommand) Handler(_ context.Context, req *admin.CommandRequest) (interface{}, error) {
    33  	input, ok := req.Data.(map[string]interface{})
    34  	if !ok {
    35  		return admin.NewInvalidAdminReqFormatError("expected { \"command\": \"add|remove|get|get_config|set_config\", \"addresses\": \"addresses\""), nil
    36  	}
    37  
    38  	command, ok := input["command"]
    39  	if !ok {
    40  		return admin.NewInvalidAdminReqErrorf("the \"command\" field is empty, must be one of add|remove|get|get_config|set_config"), nil
    41  	}
    42  
    43  	cmd, ok := command.(string)
    44  	if !ok {
    45  		return admin.NewInvalidAdminReqErrorf("the \"command\" field is not string, must be one of add|remove|get|get_config|set_config"), nil
    46  	}
    47  
    48  	if cmd == "get" {
    49  		list := s.limiter.GetAddresses()
    50  		return fmt.Sprintf("rate limited list contains a total of %d addresses: %v", len(list), list), nil
    51  	}
    52  
    53  	if cmd == "add" || cmd == "remove" {
    54  		result, ok := input["addresses"]
    55  		if !ok {
    56  			return admin.NewInvalidAdminReqErrorf("the \"addresses\" field is empty, must be hex formated addresses, can be splitted by \",\""), nil
    57  		}
    58  		addresses, ok := result.(string)
    59  		if !ok {
    60  			return admin.NewInvalidAdminReqErrorf("the \"addresses\" field is not string, must be hex formated addresses, can be splitted by \",\""), nil
    61  		}
    62  
    63  		log.Info().Msgf("admintool %v addresses: %v", cmd, addresses)
    64  
    65  		resp, err := s.AddOrRemove(cmd, addresses)
    66  		if err != nil {
    67  			return nil, err
    68  		}
    69  		return resp, nil
    70  	}
    71  
    72  	if cmd == "get_config" {
    73  		limit, burst := s.limiter.GetLimitConfig()
    74  		return fmt.Sprintf("limit: %v, burst: %v", limit, burst), nil
    75  	}
    76  
    77  	if cmd == "set_config" {
    78  		dataLimit, limit_ok := input["limit"]
    79  		dataBurst, burst_ok := input["burst"]
    80  		if !burst_ok || !limit_ok {
    81  			return admin.NewInvalidAdminReqErrorf("the \"limit\" or \"burst\" field is empty, must be number"), nil
    82  		}
    83  		limit, ok := dataLimit.(float64)
    84  		if !ok {
    85  			return admin.NewInvalidAdminReqErrorf("the \"limit\" field is not number: %v", dataLimit), nil
    86  		}
    87  
    88  		burst, ok := dataBurst.(float64)
    89  		if !ok {
    90  			return admin.NewInvalidAdminReqErrorf("the \"burst\" field is not number: %v", dataBurst), nil
    91  		}
    92  
    93  		oldLimit, oldBurst := s.limiter.GetLimitConfig()
    94  		log.Info().Msgf("admintool set_config limit: %v, burst: %v, old limit: %v, old burst: %v", limit, burst, oldLimit, oldBurst)
    95  		s.limiter.SetLimitConfig(rate.Limit(limit), int(burst))
    96  		return fmt.Sprintf("succesfully set limit %v, burst %v", limit, burst), nil
    97  	}
    98  
    99  	return fmt.Sprintf(
   100  		"invalid command field (%s), must be either \"add\" or \"remove\" or \"get\" or \"get_config\" or \"set_config\"",
   101  		cmd), nil
   102  }
   103  
   104  func (s *TxRateLimitCommand) Validator(req *admin.CommandRequest) error {
   105  	return nil
   106  }
   107  
   108  func (s *TxRateLimitCommand) AddOrRemove(command string, addresses string) (string, error) {
   109  	addrList, err := ingest.ParseAddresses(addresses)
   110  	if err != nil {
   111  		return "", err
   112  	}
   113  
   114  	if command == "add" {
   115  		ingest.AddAddresses(s.limiter, addrList)
   116  		return fmt.Sprintf("added %d addresses", len(addrList)), nil
   117  	}
   118  
   119  	// command == "remove"
   120  	ingest.RemoveAddresses(s.limiter, addrList)
   121  	return fmt.Sprintf("removed %d addresses", len(addrList)), nil
   122  }