<commit_before>package command

import (
	"flag"
	"fmt"
	"log"
	"strings"
	"sync"
	"time"

	_ "github.com/lib/pq"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/kelseyhightower/envconfig"
	"github.com/nabeken/aws-go-sqs/queue"
	"github.com/nabeken/aws-go-sqs/queue/option"
	"github.com/nabeken/delayd2"
)

type BenchSendConfig struct {
	QueueName       string `envconfig:"queue_name"`
	TargetQueueName string `envconfig:"target_queue_name"`
	Region          string `envconfig:"region"`

	NumberOfMessages int
	Concurrency      int
}

type BenchRecvConfig struct {
	QueueName string `envconfig:"queue_name"`
	Region    string `envconfig:"region"`

	NumberOfMessages int
}

type BenchCommand struct {
	Meta

	ShutdownCh <-chan struct{}
}

func (c *BenchCommand) Recv(args []string) int {
	var config BenchRecvConfig

	if err := envconfig.Process("delayd2", &config); err != nil {
		c.Ui.Error(fmt.Sprintf("Unable to get configuration from envvars: %s", err))
		return 1
	}

	cmdFlags := flag.NewFlagSet("batch recv", flag.ContinueOnError)
	cmdFlags.IntVar(&config.NumberOfMessages, "n", 1000, "number of messages")

	if err := cmdFlags.Parse(args); err != nil {
		c.Ui.Error(c.Help())
		return 1
	}

	sqsSvc := sqs.New(&aws.Config{Region: aws.String(config.Region)})
	q, err := queue.New(sqsSvc, config.QueueName)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Unable to initialize SQS connection: %s", err))
		return 1
	}

	log.Print("bench: --- Configuration ---")
	log.Printf("bench: Number of messages to receive: %d", config.NumberOfMessages)
	log.Printf("bench: DELAYD2_QUEUE_NAME=%s", config.QueueName)

	errCh := make(chan error)

	log.Print("bench: starting")
	go func() {
		payloads := map[string]struct{}{}
		for {
			log.Printf("%d messages processed", len(payloads))
			if len(payloads) >= config.NumberOfMessages {
				break
			}

			messages, err := q.ReceiveMessage(
				option.MaxNumberOfMessages(10),
				option.UseAllAttribute(),
			)
			// continue even if we get an error
			if err != nil {
				log.Printf("unable to receive messages but continuing...: %s", err)
				continue
			}

			recipientHandles := make([]*string, 0, 10)
			for _, m := range messages {
				payloads[*m.Body] = struct{}{}
				recipientHandles = append(recipientHandles, m.ReceiptHandle)
			}

			if err := q.DeleteMessageBatch(recipientHandles...); err != nil {
				log.Printf("unable to delete message but continuing...: %s", err)
			}
		}

		log.Printf("%d messages received", config.NumberOfMessages)
		errCh <- nil
	}()

	select {
	case <-errCh:
		if err != nil {
			c.Ui.Error(fmt.Sprintf("Unable to receive messages: %s", err))
			return 1
		}
		log.Print("bench: done.")
	case <-c.ShutdownCh:
		log.Fatal("signal received.")
	}

	return 0
}

func (c *BenchCommand) Send(args []string) int {
	var config BenchSendConfig

	if err := envconfig.Process("delayd2", &config); err != nil {
		c.Ui.Error(fmt.Sprintf("Unable to get configuration from envvars: %s", err))
		return 1
	}

	cmdFlags := flag.NewFlagSet("batch send", flag.ContinueOnError)
	cmdFlags.IntVar(&config.NumberOfMessages, "n", 1000, "number of messages")
	cmdFlags.IntVar(&config.Concurrency, "c", 5, "concurrency")

	if err := cmdFlags.Parse(args); err != nil {
		c.Ui.Error(c.Help())
		return 1
	}

	sqsSvc := sqs.New(&aws.Config{Region: aws.String(config.Region)})
	q, err := queue.New(sqsSvc, config.QueueName)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Unable to initialize SQS connection: %s", err))
		return 1
	}

	log.Print("bench: --- Configuration ---")
	log.Printf("bench: Number of messages to send: %d", config.NumberOfMessages)
	log.Printf("bench: Concurrency of sending: %d", config.Concurrency)
	log.Printf("bench: DELAYD2_QUEUE_NAME=%s", config.QueueName)
	log.Printf("bench: DELAYD2_TARGET_QUEUE_NAME=%s", config.TargetQueueName)

	s := delayd2.NewSender(q)

	log.Print("bench: starting")
	payloads := make([]string, config.NumberOfMessages)
	for i := 0; i < config.NumberOfMessages; i++ {
		payloads[i] = fmt.Sprintf("%d %d", i, time.Now().Unix())
	}

	tasks := DistributeN(config.Concurrency, payloads)

	var wg sync.WaitGroup
	wg.Add(len(tasks))

	begin := time.Now()
	for _, task := range tasks {
		go func(t []string) {
			defer wg.Done()

			log.Print("launching worker goroutine")
			if err := s.SendMessageBatch(0, config.TargetQueueName, t); err != nil {
				log.Print(err)
				return
			}
		}(task)
	}

	doneCh := make(chan struct{})
	go func() {<commit_msg>bench: Add duration support for bench send<commit_after>package command

import (
	"flag"
	"fmt"
	"log"
	"strings"
	"sync"
	"time"

	_ "github.com/lib/pq"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/kelseyhightower/envconfig"
	"github.com/nabeken/aws-go-sqs/queue"
	"github.com/nabeken/aws-go-sqs/queue/option"
	"github.com/nabeken/delayd2"
)

type BenchRecvConfig struct {
	QueueName string `envconfig:"queue_name"`
	Region    string `envconfig:"region"`

	NumberOfMessages int
}

type BenchCommand struct {
	Meta

	ShutdownCh <-chan struct{}
}

func (c *BenchCommand) Recv(args []string) int {
	var config BenchRecvConfig

	if err := envconfig.Process("delayd2", &config); err != nil {
		c.Ui.Error(fmt.Sprintf("Unable to get configuration from envvars: %s", err))
		return 1
	}

	cmdFlags := flag.NewFlagSet("batch recv", flag.ContinueOnError)
	cmdFlags.IntVar(&config.NumberOfMessages, "n", 1000, "number of messages")

	if err := cmdFlags.Parse(args); err != nil {
		c.Ui.Error(c.Help())
		return 1
	}

	sqsSvc := sqs.New(&aws.Config{Region: aws.String(config.Region)})
	q, err := queue.New(sqsSvc, config.QueueName)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Unable to initialize SQS connection: %s", err))
		return 1
	}

	log.Print("bench: --- Configuration ---")
	log.Printf("bench: Number of messages to receive: %d", config.NumberOfMessages)
	log.Printf("bench: DELAYD2_QUEUE_NAME=%s", config.QueueName)

	errCh := make(chan error)

	log.Print("bench: starting")
	go func() {
		payloads := map[string]struct{}{}
		for {
			log.Printf("%d messages processed", len(payloads))
			if len(payloads) >= config.NumberOfMessages {
				break
			}

			messages, err := q.ReceiveMessage(
				option.MaxNumberOfMessages(10),
				option.UseAllAttribute(),
			)
			// continue even if we get an error
			if err != nil {
				log.Printf("unable to receive messages but continuing...: %s", err)
				continue
			}

			recipientHandles := make([]*string, 0, 10)
			for _, m := range messages {
				payloads[*m.Body] = struct{}{}
				recipientHandles = append(recipientHandles, m.ReceiptHandle)
			}

			if err := q.DeleteMessageBatch(recipientHandles...); err != nil {
				log.Printf("unable to delete message but continuing...: %s", err)
			}
		}

		log.Printf("%d messages received", config.NumberOfMessages)
		errCh <- nil
	}()

	select {
	case <-errCh:
		if err != nil {
			c.Ui.Error(fmt.Sprintf("Unable to receive messages: %s", err))
			return 1
		}
		log.Print("bench: done.")
	case <-c.ShutdownCh:
		log.Fatal("signal received.")
	}

	return 0
}

type BenchSendConfig struct {
	QueueName       string `envconfig:"queue_name"`
	TargetQueueName string `envconfig:"target_queue_name"`
	Region          string `envconfig:"region"`

	NumberOfMessages int
	Concurrency      int
	Duration         int
}

func (c *BenchCommand) Send(args []string) int {
	var config BenchSendConfig

	if err := envconfig.Process("delayd2", &config); err != nil {
		c.Ui.Error(fmt.Sprintf("Unable to get configuration from envvars: %s", err))
		return 1
	}

	cmdFlags := flag.NewFlagSet("batch send", flag.ContinueOnError)
	cmdFlags.IntVar(&config.NumberOfMessages, "n", 1000, "number of messages")
	cmdFlags.IntVar(&config.Concurrency, "c", 5, "concurrency")
	cmdFlags.IntVar(&config.Duration, "d", 10, "duration in second")

	if err := cmdFlags.Parse(args); err != nil {
		c.Ui.Error(c.Help())
		return 1
	}

	sqsSvc := sqs.New(&aws.Config{Region: aws.String(config.Region)})
	q, err := queue.New(sqsSvc, config.QueueName)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Unable to initialize SQS connection: %s", err))
		return 1
	}

	log.Print("bench: --- Configuration ---")
	log.Printf("bench: Number of messages to send: %d", config.NumberOfMessages)
	log.Printf("bench: Concurrency of sending: %d", config.Concurrency)
	log.Printf("bench: Duration of relaying: %d", config.Duration)
	log.Printf("bench: DELAYD2_QUEUE_NAME=%s", config.QueueName)
	log.Printf("bench: DELAYD2_TARGET_QUEUE_NAME=%s", config.TargetQueueName)

	s := delayd2.NewSender(q)

	log.Print("bench: starting")
	payloads := make([]string, config.NumberOfMessages)
	for i := 0; i < config.NumberOfMessages; i++ {
		payloads[i] = fmt.Sprintf("%d %d", i, time.Now().Unix())
	}

	tasks := DistributeN(config.Concurrency, payloads)

	var wg sync.WaitGroup
	wg.Add(len(tasks))

	begin := time.Now()
	for _, task := range tasks {
		go func(t []string) {
			defer wg.Done()

			log.Print("launching worker goroutine")
			if err := s.SendMessageBatch(config.Duration, config.TargetQueueName, t); err != nil {
				log.Print(err)
				return
			}
		}(task)
	}

	doneCh := make(chan struct{})
	go func() {
{'fraction_non_alphanumeric': 0.15559192580469175, 'fraction_numerical': 0.006873977086743044, 'mean_word_length': 4.016967706622879, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '637566', 'n_tokens_mistral': 3546, 'n_tokens_neox': 3231, 'n_words': 746}