<commit_before>//   Copyright 2013 <NAME> <<EMAIL>>
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.

package main

import (
	"flag"
	"github.com/kr/beanstalk"
	"log"
	"time"
)

// Get Parameters from cli
var publishers = flag.Int("p", 1, "number of concurrent publishers, default to 1")
var readers = flag.Int("r", *publishers, "number of concurrent readers, default to number of publishers")
var count = flag.Int("n", 10000, "Count of jobs to be processed, default to 10000")
var host = flag.String("h", "localhost:11300", "Host to beanstalkd, default to localhost:11300")
var size = flag.Int("s", 256, "Size of data, default to 256. in byte")
var drain = flag.Bool("d", false, "Drain the beanstalk before stating test")

func testPublisher(h string, count int, size int, ch chan int) {
	conn, e := beanstalk.Dial("tcp", h)
	defer conn.Close()
	data := make([]byte, size)
	if e != nil {
		log.Fatal(e)
	}
	for i := 0; i < count; i++ {
		_, err := conn.Put(data, 0, 0, 120*time.Second)
		if err != nil {
			log.Fatal(err)
		}
	}
	ch <- 1
}

func testReader(h string, count int, ch chan int) {
	conn, e := beanstalk.Dial("tcp", h)
	defer conn.Close()
	if e != nil {
		log.Fatal(e)
	}
	for i := 0; i < count; i++ {
		id, _, e := conn.Reserve(250 * time.Millisecond)
		if e != nil {
			log.Println(e)
			continue
		}
		e = conn.Delete(id)
		if e != nil {
			log.Println(e)
		}
	}
	ch <- 1
}

func drainBeanstalk(h string) {
	log.Println("Draining beanstalk")
	conn, e := beanstalk.Dial("tcp", h)
	defer conn.Close()
	if e != nil {
		log.Fatal(e)
	}
	for {
		id, _, e := conn.Reserve(250 * time.Millisecond)
		if e != nil {
			return
		}
		e = conn.Delete(id)
		if e != nil {
			log.Println(e)
		}
	}
}

func main() {
	flag.Parse()
	if *drain {
		drainBeanstalk(*host)
	}

	log.Println("Target host: ", *host)
	log.Println("Starting publishers: ", *publishers)
	log.Println("Starting readers: ", *readers)
	log.Println("Total jobs to be processed: ", *count)
	log.Println("Benchmarking, be patient ...")
	chPublisher := make(chan int)
	chReader := make(chan int)
	t0 := time.Now()

	publishCount := *count / *publishers
	for i := 0; i < *publishers; i++ {
		go testPublisher(*host, publishCount, *size, chPublisher)
	}

	readCount := *count / *readers
	for i := 0; i < *readers; i++ {
		go testReader(*host, readCount, chReader)
	}

	// Wait for return, assume publishers will finish first
	for i := 0; i < *publishers; i++ {
		<-chPublisher
	}

	log.Println("---------------")
	delta := time.Now().Sub(t0)
	log.Println("Publishers finished at: ", delta)
	log.Println("Publish rate: ", float64(*count)/delta.Seconds(), " req/s")

	for i := 0; i < *readers; i++ {
		<-chReader
	}
	delta = time.Now().Sub(t0)
	log.Println("Readers finished at: ", delta)
	log.Println("Read rate: ", float64(*count)/delta.Seconds(), " req/s")
}
<commit_msg>Add option to fill beanstalk to a particular number of jobs before starting benchmark<commit_after>//   Copyright 2013 <NAME> <<EMAIL>>
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.

package main

import (
	"flag"
	"github.com/kr/beanstalk"
	"log"
	"time"
)

// Get Parameters from cli
var publishers = flag.Int("p", 1, "number of concurrent publishers, default to 1")
var readers = flag.Int("r", *publishers, "number of concurrent readers, default to number of publishers")
var count = flag.Int("n", 10000, "Count of jobs to be processed, default to 10000")
var host = flag.String("h", "localhost:11300", "Host to beanstalkd, default to localhost:11300")
var size = flag.Int("s", 256, "Size of data, default to 256. in byte")
var drain = flag.Bool("d", false, "Drain the beanstalk before starting test")
var fill = flag.Int("f", 0, "Place <f> jobs on the beanstalk before starting test")

func testPublisher(h string, count int, size int, ch chan int) {
	conn, e := beanstalk.Dial("tcp", h)
	defer conn.Close()
	data := make([]byte, size)
	if e != nil {
		log.Fatal(e)
	}
	for i := 0; i < count; i++ {
		_, err := conn.Put(data, 0, 0, 120*time.Second)
		if err != nil {
			log.Fatal(err)
		}
	}
	ch <- 1
}

func testReader(h string, count int, ch chan int) {
	conn, e := beanstalk.Dial("tcp", h)
	defer conn.Close()
	if e != nil {
		log.Fatal(e)
	}
	for i := 0; i < count; i++ {
		id, _, e := conn.Reserve(250 * time.Millisecond)
		if e != nil {
			log.Println(e)
			continue
		}
		e = conn.Delete(id)
		if e != nil {
			log.Println(e)
		}
	}
	ch <- 1
}

func drainBeanstalk(h string) {
	log.Println("Draining beanstalk")
	conn, e := beanstalk.Dial("tcp", h)
	defer conn.Close()
	if e != nil {
		log.Fatal(e)
	}
	for {
		id, _, e := conn.Reserve(250 * time.Millisecond)
		if e != nil {
			return
		}
		e = conn.Delete(id)
		if e != nil {
			log.Println(e)
		}
	}
}

func fillBeanstalk(h string, count int, size int) {
	log.Println("Filling beanstalk")
	ch := make(chan int)
	go testPublisher(h, count, size, ch)
	<-ch
}

func main() {
	flag.Parse()
	if *drain {
		drainBeanstalk(*host)
	}
	if (*fill) > 0 {
		fillBeanstalk(*host, *fill, *size)
	}

	log.Println("Target host: ", *host)
	log.Println("Starting publishers: ", *publishers)
	log.Println("Starting readers: ", *readers)
	log.Println("Total jobs to be processed: ", *count)
	log.Println("Benchmarking, be patient ...")
	chPublisher := make(chan int)
	chReader := make(chan int)
	t0 := time.Now()

	publishCount := *count / *publishers
	for i := 0; i < *publishers; i++ {
		go testPublisher(*host, publishCount, *size, chPublisher)
	}

	readCount := *count / *readers
	for i := 0; i < *readers; i++ {
		go testReader(*host, readCount, chReader)
	}

	// Wait for return, assume publishers will finish first
	for i := 0; i < *publishers; i++ {
		<-chPublisher
	}

	log.Println("---------------")
	delta := time.Now().Sub(t0)
	log.Println("Publishers finished at: ", delta)
	log.Println("Publish rate: ", float64(*count)/delta.Seconds(), " req/s")

	for i := 0; i < *readers; i++ {
		<-chReader
	}
	delta = time.Now().Sub(t0)
	log.Println("Readers finished at: ", delta)
	log.Println("Read rate: ", float64(*count)/delta.Seconds(), " req/s")
}

{'fraction_non_alphanumeric': 0.16050254093732355, 'fraction_numerical': 0.017786561264822136, 'mean_word_length': 3.422596754057428, 'pattern_counts': {'":': 0, '<': 31, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, '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': '4589770', 'n_tokens_mistral': 2751, 'n_tokens_neox': 2486, 'n_words': 802}