Merge pull request #13 from v3io/restore

Restore

httpblaster/executor.go

@@ -76,6 +76,12 @@ func (self *Executor) load_request_generator() (chan *fasthttp.Request, bool) {
  case request_generators.JSON2KV:
  req_gen = &request_generators.Json2KV{}
  break
+ case request_generators.LINE2KV:
+ req_gen = &request_generators.Line2KvGenerator{}
+ break
+ case request_generators.RESTORE:
+ req_gen = &request_generators.RestoreGenerator{}
+ break
  default:
  panic(fmt.Sprintf("unknown request generator %s", self.Workload.Generator))
  }

httpblaster/request_generators/common.go

@@ -13,31 +13,52 @@ const (
  LINE2STREAM = "line2stream"
  CSV2KV = "csv2kv"
  JSON2KV = "json2kv"
+ LINE2KV = "line2kv"
+ RESTORE = "restore"
 )
 
 type RequestCommon struct {
  ch_files chan string
  base_uri string
 }
 
+var (
+ contentType string = "application/json"
+)
+
 func (self *RequestCommon) PrepareRequest(content_type string,
  header_args map[string]string,
  method string, uri string,
  body string, host string) *fasthttp.Request {
  req := fasthttp.AcquireRequest()
 
- header := fasthttp.RequestHeader{}
- header.SetContentType(content_type)
+ req.Header.SetContentType(content_type)
+ req.Header.SetMethod(method)
+ req.Header.SetRequestURI(uri)
+ req.Header.SetHost(host)
+ for k, v := range header_args {
+ req.Header.Set(k, v)
+ }
 
- header.SetMethod(method)
- header.SetRequestURI(uri)
- header.SetHost(host)
+ req.AppendBodyString(body)
+ return req
+}
 
+func (self *RequestCommon) PrepareRequestBytes(content_type string,
+ header_args map[string]string,
+ method string, uri string,
+ body []byte, host string) *fasthttp.Request {
+ req := fasthttp.AcquireRequest()
+
+ req.Header.SetContentType(content_type)
+ req.Header.SetMethod(method)
+ req.Header.SetRequestURI(uri)
+ req.Header.SetHost(host)
  for k, v := range header_args {
- header.Set(k, v)
+ req.Header.Set(k, v)
  }
- req.AppendBodyString(body)
- header.CopyTo(&req.Header)
+
+ req.AppendBody(body)
  return req
 }
 

httpblaster/request_generators/line2kv.go

@@ -0,0 +1,94 @@
+package request_generators
+
+import (
+ "bufio"
+ "fmt"
+ //"github.com/nu7hatch/gouuid"
+ "github.com/v3io/http_blaster/httpblaster/config"
+ //"github.com/v3io/http_blaster/httpblaster/igz_data"
+ "github.com/valyala/fasthttp"
+ "io"
+ "log"
+ "os"
+ "runtime"
+ "strings"
+ "sync"
+)
+
+type Line2KvGenerator struct {
+ RequestCommon
+ workload config.Workload
+}
+
+func (self *Line2KvGenerator) UseCommon(c RequestCommon) {
+
+}
+
+func (self *Line2KvGenerator) generate_request(ch_records chan []string,
+ ch_req chan *fasthttp.Request,
+ host string, wg *sync.WaitGroup) {
+ defer wg.Done()
+ for r := range ch_records {
+ req := self.PrepareRequest(contentType, self.workload.Header, "PUT",
+ r[0], r[1], host)
+ //panic(fmt.Sprintf("%+v",r))
+ ch_req <- req
+ }
+ log.Println("generate_request Done")
+}
+
+func (self *Line2KvGenerator) generate(ch_req chan *fasthttp.Request, payload string, host string) {
+ defer close(ch_req)
+ var ch_records chan []string = make(chan []string)
+ wg := sync.WaitGroup{}
+ ch_files := self.FilesScan(self.workload.Payload)
+
+ wg.Add(runtime.NumCPU())
+ for c := 0; c < runtime.NumCPU(); c++ {
+ go self.generate_request(ch_records, ch_req, host, &wg)
+ }
+
+ for f := range ch_files {
+ if file, err := os.Open(f); err == nil {
+ reader := bufio.NewReader(file)
+ var i int = 0
+ for {
+ address, addr_err := reader.ReadString('\n')
+ payload, payload_err := reader.ReadString('\n')
+
+ if addr_err == nil && payload_err == nil {
+ ch_records <- []string{strings.TrimSpace(address), string(payload)}
+ i++
+ } else if addr_err == io.EOF || payload_err == io.EOF {
+ break
+ } else {
+ log.Fatal(err)
+ }
+ }
+
+ log.Println(fmt.Sprintf("Finish file scaning, generated %d records", i))
+ } else {
+ panic(err)
+ }
+ }
+ close(ch_records)
+ log.Println("Waiting for generators to finish")
+ wg.Wait()
+ log.Println("generators done")
+}
+
+func (self *Line2KvGenerator) GenerateRequests(wl config.Workload, tls_mode bool, host string) chan *fasthttp.Request {
+ self.workload = wl
+ if self.workload.Header == nil {
+ self.workload.Header = make(map[string]string)
+ }
+ self.workload.Header["X-v3io-function"] = "PutItem"
+
+ self.SetBaseUri(tls_mode, host, self.workload.Container, self.workload.Target)
+
+ ch_req := make(chan *fasthttp.Request, 1000)
+
+ go self.generate(ch_req, self.workload.Payload, host)
+
+ return ch_req
+}

httpblaster/request_generators/line2stream.go

@@ -33,7 +33,7 @@ func (self *Line2StreamGenerator) generate_request(ch_records chan string,
  for r := range ch_records {
  sr := igz_data.NewStreamRecord("client", r, u.String(), 0, true)
  r := igz_data.NewStreamRecords(sr)
- req := self.PrepareRequest(contentType, self.workload.Header, self.workload.Type,
+ req := self.PrepareRequest(contentType, self.workload.Header, "PUT",
  self.base_uri, r.ToJsonString(), host)
  ch_req <- req
  }