<commit_before>package logger

import (
	"fmt"
	"log"
	"sort"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatchlogs"
	"github.com/honeybadger-io/honeybadger-go"
	"github.com/satori/go.uuid"
)

const (
	maxBatchByteSize = 1048576
	maxBatchLength   = 10000
	logEventOverhead = 26
)

type Logger interface {
	Log(t time.Time, s string)
}

type logBatch []*cloudwatchlogs.InputLogEvent

func (b logBatch) Len() int {
	return len(b)
}

func (b logBatch) Less(i, j int) bool {
	return *b[i].Timestamp < *b[j].Timestamp
}

func (b logBatch) Swap(i, j int) {
	b[i], b[j] = b[j], b[i]
}

type CloudWatchLogger struct {
	logGroupName  string
	logStreamName string
	sequenceToken *string
	retention     int
	logs          chan *cloudwatchlogs.InputLogEvent
	batch         logBatch
	batchByteSize int
	timeout       <-chan time.Time
	client        *cloudwatchlogs.CloudWatchLogs
}

func NewCloudWatchLogger(logGroupName string, retention int) (*CloudWatchLogger, error) {
	sess, err := session.NewSession()
	if err != nil {
		return nil, fmt.Errorf("failed to create AWS session: %s", err)
	}

	client := cloudwatchlogs.New(sess)

	cwl := &CloudWatchLogger{
		logGroupName:  logGroupName,
		logStreamName: uuid.NewV4().String(),
		retention:     retention,
		logs:          make(chan *cloudwatchlogs.InputLogEvent, 100),
		client:        client,
	}
	go cwl.worker()
	return cwl, nil
}

func (cwl *CloudWatchLogger) Log(t time.Time, s string) {
	cwl.logs <- &cloudwatchlogs.InputLogEvent{
		Message:   aws.String(s),
		Timestamp: aws.Int64(t.UnixNano() / int64(time.Millisecond)),
	}
}

func (cwl *CloudWatchLogger) worker() {
	cwl.resetBatch()
	for {
		select {
		case logEvent := <-cwl.logs:
			cwl.addToBatch(logEvent)
		case <-cwl.timeout:
			cwl.flush()
		}
	}
}

func (cwl *CloudWatchLogger) addToBatch(logEvent *cloudwatchlogs.InputLogEvent) {
	logEventSize := len(*logEvent.Message) + logEventOverhead

	if logEventSize+cwl.batchByteSize > maxBatchByteSize || len(cwl.batch) == maxBatchLength {
		cwl.flush()
	}

	if cwl.timeout == nil {
		cwl.timeout = time.After(time.Second)
	}

	cwl.batch = append(cwl.batch, logEvent)
	cwl.batchByteSize += logEventSize
}

func (cwl *CloudWatchLogger) flush() {
	batch := cwl.batch
	cwl.resetBatch()
	sort.Sort(batch)
	if err := cwl.sendToCloudWatchLogs(batch); err != nil {
		if honeybadger.Config.APIKey == "" {
			honeybadger.Notify(err)
		}
		log.Println(err)
	}
}

func (cwl *CloudWatchLogger) resetBatch() {
	cwl.batch = logBatch{}
	cwl.batchByteSize = 0
	cwl.timeout = nil
}

func (cwl *CloudWatchLogger) sendToCloudWatchLogs(batch logBatch) error {
	s := time.Now()
	params := &cloudwatchlogs.PutLogEventsInput{
		LogEvents:     batch,
		LogGroupName:  aws.String(cwl.logGroupName),
		LogStreamName: aws.String(cwl.logStreamName),
		SequenceToken: cwl.sequenceToken,
	}
	resp, err := cwl.client.PutLogEvents(params)

	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ResourceNotFoundException" {
				if err = cwl.createLogStream(); err != nil {
					return err
				}
				return cwl.sendToCloudWatchLogs(batch)
			}
		}
		return fmt.Errorf("PutLogEvents failed: %s", err)
	}
	log.Printf("wrote %d log events (%d bytes) in %s\n", len(batch), cwl.batchByteSize, time.Since(s))

	cwl.sequenceToken = resp.NextSequenceToken
	return nil
}

func (cwl *CloudWatchLogger) createLogStream() error {
	params := &cloudwatchlogs.CreateLogStreamInput{
		LogGroupName:  aws.String(cwl.logGroupName),
		LogStreamName: aws.String(cwl.logStreamName),<commit_msg>Pass the batch size to sendToCloudWatchLogs<commit_after>package logger

import (
	"fmt"
	"log"
	"sort"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatchlogs"
	"github.com/honeybadger-io/honeybadger-go"
	"github.com/satori/go.uuid"
)

const (
	maxBatchByteSize = 1048576
	maxBatchLength   = 10000
	logEventOverhead = 26
)

type Logger interface {
	Log(t time.Time, s string)
}

type logBatch []*cloudwatchlogs.InputLogEvent

func (b logBatch) Len() int {
	return len(b)
}

func (b logBatch) Less(i, j int) bool {
	return *b[i].Timestamp < *b[j].Timestamp
}

func (b logBatch) Swap(i, j int) {
	b[i], b[j] = b[j], b[i]
}

type CloudWatchLogger struct {
	logGroupName  string
	logStreamName string
	sequenceToken *string
	retention     int
	logs          chan *cloudwatchlogs.InputLogEvent
	batch         logBatch
	batchByteSize int
	timeout       <-chan time.Time
	client        *cloudwatchlogs.CloudWatchLogs
}

func NewCloudWatchLogger(logGroupName string, retention int) (*CloudWatchLogger, error) {
	sess, err := session.NewSession()
	if err != nil {
		return nil, fmt.Errorf("failed to create AWS session: %s", err)
	}

	client := cloudwatchlogs.New(sess)

	cwl := &CloudWatchLogger{
		logGroupName:  logGroupName,
		logStreamName: uuid.NewV4().String(),
		retention:     retention,
		logs:          make(chan *cloudwatchlogs.InputLogEvent, 100),
		client:        client,
	}
	go cwl.worker()
	return cwl, nil
}

func (cwl *CloudWatchLogger) Log(t time.Time, s string) {
	cwl.logs <- &cloudwatchlogs.InputLogEvent{
		Message:   aws.String(s),
		Timestamp: aws.Int64(t.UnixNano() / int64(time.Millisecond)),
	}
}

func (cwl *CloudWatchLogger) worker() {
	cwl.resetBatch()
	for {
		select {
		case logEvent := <-cwl.logs:
			cwl.addToBatch(logEvent)
		case <-cwl.timeout:
			cwl.flush()
		}
	}
}

func (cwl *CloudWatchLogger) addToBatch(logEvent *cloudwatchlogs.InputLogEvent) {
	logEventSize := len(*logEvent.Message) + logEventOverhead

	if logEventSize+cwl.batchByteSize > maxBatchByteSize || len(cwl.batch) == maxBatchLength {
		cwl.flush()
	}

	if cwl.timeout == nil {
		cwl.timeout = time.After(time.Second)
	}

	cwl.batch = append(cwl.batch, logEvent)
	cwl.batchByteSize += logEventSize
}

func (cwl *CloudWatchLogger) flush() {
	batch := cwl.batch
	batchByteSize := cwl.batchByteSize
	cwl.resetBatch()
	sort.Sort(batch)
	if err := cwl.sendToCloudWatchLogs(batch, batchByteSize); err != nil {
		if honeybadger.Config.APIKey == "" {
			honeybadger.Notify(err)
		}
		log.Println(err)
	}
}

func (cwl *CloudWatchLogger) resetBatch() {
	cwl.batch = logBatch{}
	cwl.batchByteSize = 0
	cwl.timeout = nil
}

func (cwl *CloudWatchLogger) sendToCloudWatchLogs(batch logBatch, batchByteSize int) error {
	s := time.Now()
	params := &cloudwatchlogs.PutLogEventsInput{
		LogEvents:     batch,
		LogGroupName:  aws.String(cwl.logGroupName),
		LogStreamName: aws.String(cwl.logStreamName),
		SequenceToken: cwl.sequenceToken,
	}
	resp, err := cwl.client.PutLogEvents(params)

	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ResourceNotFoundException" {
				if err = cwl.createLogStream(); err != nil {
					return err
				}
				return cwl.sendToCloudWatchLogs(batch, batchByteSize)
			}
		}
		return fmt.Errorf("PutLogEvents failed: %s", err)
	}
	log.Printf("wrote %d log events (%d bytes) in %s\n", len(batch), batchByteSize, time.Since(s))

	cwl.sequenceToken = resp.NextSequenceToken
	return nil
}

func (cwl *CloudWatchLogger) createLogStream() error {
	params := &cloudwatchlogs.CreateLogStreamInput{
		LogGroupName:  aws.String(cwl.logGroupName),
		LogStreamName: aws.String(cwl.logStreamName),
{'fraction_non_alphanumeric': 0.13354079285617643, 'fraction_numerical': 0.006223785685292924, 'mean_word_length': 4.083906464924347, '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': '1599017', 'n_tokens_mistral': 2891, 'n_tokens_neox': 2738, 'n_words': 558}