<commit_before>package lxd

import (
	"fmt"
	"log"

	"github.com/hashicorp/terraform/helper/schema"
)

func resourceLxdContainerFile() *schema.Resource {
	return &schema.Resource{
		Create: resourceLxdContainerFileCreate,
		Exists: resourceLxdContainerFileExists,
		Read:   resourceLxdContainerFileRead,
		Delete: resourceLxdContainerFileDelete,

		Schema: map[string]*schema.Schema{
			"container_name": &schema.Schema{
				Type:     schema.TypeString,
				ForceNew: true,
				Required: true,
			},

			"target_file": &schema.Schema{
				Type:     schema.TypeString,
				ForceNew: true,
				Required: true,
			},

			"content": &schema.Schema{
				Type:          schema.TypeString,
				ForceNew:      true,
				Optional:      true,
				ConflictsWith: []string{"source"},
			},

			"source": &schema.Schema{
				Type:          schema.TypeString,
				ForceNew:      true,
				Optional:      true,
				ConflictsWith: []string{"content"},
			},

			"uid": &schema.Schema{
				Type:     schema.TypeInt,
				ForceNew: true,
				Optional: true,
			},

			"gid": &schema.Schema{
				Type:     schema.TypeInt,
				ForceNew: true,
				Optional: true,
			},

			"mode": &schema.Schema{
				Type:     schema.TypeString,
				ForceNew: true,
				Optional: true,
			},

			"create_directories": &schema.Schema{
				Type:     schema.TypeBool,
				ForceNew: true,
				Optional: true,
			},

			"remote": &schema.Schema{
				Type:     schema.TypeString,
				ForceNew: true,
				Optional: true,
			},
		},
	}
}

func resourceLxdContainerFileCreate(d *schema.ResourceData, meta interface{}) error {
	p := meta.(*lxdProvider)
	remote := p.selectRemote(d)
	server, err := p.GetContainerServer(remote)
	if err != nil {
		return err
	}

	containerName := d.Get("container_name").(string)
	_, _, err = server.GetContainer(containerName)
	if err != nil {
		return fmt.Errorf("unable to retrieve container %s: %s", containerName, err)
	}

	file := File{
		RemoteName:        remote,
		ContainerName:     containerName,
		TargetFile:        d.Get("target_file").(string),
		Content:           d.Get("content").(string),
		Source:            d.Get("source").(string),
		UID:               d.Get("uid").(int),
		GID:               d.Get("gid").(int),
		Mode:              d.Get("mode").(string),
		CreateDirectories: d.Get("create_directories").(bool),
	}

	err = containerUploadFile(server, containerName, file)
	if err != nil {
		return fmt.Errorf("unable to create file %s: %s", file.TargetFile, err)
	}

	d.SetId(file.String())
	return nil
}

func resourceLxdContainerFileRead(d *schema.ResourceData, meta interface{}) error {
	p := meta.(*lxdProvider)
	v, targetFile := newFileIDFromResourceID(d.Id())
	remote, containerName, err := p.Config.ParseRemote(v)

	remote = p.selectRemote(d)
	server, err := p.GetContainerServer(remote)
	if err != nil {
		return err
	}

	_, _, err = server.GetContainer(containerName)
	if err != nil {
		return fmt.Errorf("unable to retrieve container %s: %s", containerName, err)
	}

	_, file, err := server.GetContainerFile(containerName, targetFile)
	if err != nil {
		return fmt.Errorf("unable to retrieve file %s:%s: %s", containerName, targetFile, err)
	}

	log.Printf("[DEBUG] Retrieved file: %#v", file)

	d.Set("container_name", containerName)
	d.Set("target_file", targetFile)
	d.Set("uid", file.UID)
	d.Set("gid", file.GID)
	d.Set("mode", file.Mode)

	return nil
}

func resourceLxdContainerFileDelete(d *schema.ResourceData, meta interface{}) error {
	p := meta.(*lxdProvider)
	v, targetFile := newFileIDFromResourceID(d.Id())
	remote, containerName, err := p.Config.ParseRemote(v)
	server, err := p.GetContainerServer(remote)
	if err != nil {
		return err
	}

	_, _, err = server.GetContainer(containerName)
	if err != nil {
		return fmt.Errorf("unable to retrieve container %s: %s", containerName, err)
	}

	err = containerDeleteFile(server, containerName, targetFile)
	if err != nil {
		return fmt.Errorf("unable to delete file %s:%s: %s", containerName, targetFile, err)
	}

	return nil
}

func resourceLxdContainerFileExists(d *schema.ResourceData, meta interface{}) (exists bool, err error) {
	p := meta.(*lxdProvider)
	v, targetFile := newFileIDFromResourceID(d.Id())
	remote, containerName, err := p.Config.ParseRemote(v)
	if err != nil {
		err = fmt.Errorf("unable to determine remote: %s", err)
		return
	}

	server, err := p.GetContainerServer(remote)
	if err != nil {
		return
	}

	_, _, err = server.GetContainer(containerName)
	if err != nil {
		err = fmt.Errorf("unable to retrieve container %s: %s", containerName, err)
		return
	}

	_, _, err = server.GetContainerFile(containerName, targetFile)
	if err != nil {
		return
	}

	exists = true

	return
}
<commit_msg>Container File: Recreate when not found (#130)<commit_after>package lxd

import (
	"fmt"
	"log"

	"github.com/hashicorp/terraform/helper/schema"
)

func resourceLxdContainerFile() *schema.Resource {
	return &schema.Resource{
		Create: resourceLxdContainerFileCreate,
		Exists: resourceLxdContainerFileExists,
		Read:   resourceLxdContainerFileRead,
		Delete: resourceLxdContainerFileDelete,

		Schema: map[string]*schema.Schema{
			"container_name": &schema.Schema{
				Type:     schema.TypeString,
				ForceNew: true,
				Required: true,
			},

			"target_file": &schema.Schema{
				Type:     schema.TypeString,
				ForceNew: true,
				Required: true,
			},

			"content": &schema.Schema{
				Type:          schema.TypeString,
				ForceNew:      true,
				Optional:      true,
				ConflictsWith: []string{"source"},
			},

			"source": &schema.Schema{
				Type:          schema.TypeString,
				ForceNew:      true,
				Optional:      true,
				ConflictsWith: []string{"content"},
			},

			"uid": &schema.Schema{
				Type:     schema.TypeInt,
				ForceNew: true,
				Optional: true,
			},

			"gid": &schema.Schema{
				Type:     schema.TypeInt,
				ForceNew: true,
				Optional: true,
			},

			"mode": &schema.Schema{
				Type:     schema.TypeString,
				ForceNew: true,
				Optional: true,
			},

			"create_directories": &schema.Schema{
				Type:     schema.TypeBool,
				ForceNew: true,
				Optional: true,
			},

			"remote": &schema.Schema{
				Type:     schema.TypeString,
				ForceNew: true,
				Optional: true,
			},
		},
	}
}

func resourceLxdContainerFileCreate(d *schema.ResourceData, meta interface{}) error {
	p := meta.(*lxdProvider)
	remote := p.selectRemote(d)
	server, err := p.GetContainerServer(remote)
	if err != nil {
		return err
	}

	containerName := d.Get("container_name").(string)
	_, _, err = server.GetContainer(containerName)
	if err != nil {
		return fmt.Errorf("unable to retrieve container %s: %s", containerName, err)
	}

	file := File{
		RemoteName:        remote,
		ContainerName:     containerName,
		TargetFile:        d.Get("target_file").(string),
		Content:           d.Get("content").(string),
		Source:            d.Get("source").(string),
		UID:               d.Get("uid").(int),
		GID:               d.Get("gid").(int),
		Mode:              d.Get("mode").(string),
		CreateDirectories: d.Get("create_directories").(bool),
	}

	err = containerUploadFile(server, containerName, file)
	if err != nil {
		return fmt.Errorf("unable to create file %s: %s", file.TargetFile, err)
	}

	d.SetId(file.String())
	return nil
}

func resourceLxdContainerFileRead(d *schema.ResourceData, meta interface{}) error {
	p := meta.(*lxdProvider)
	v, targetFile := newFileIDFromResourceID(d.Id())
	remote, containerName, err := p.Config.ParseRemote(v)

	remote = p.selectRemote(d)
	server, err := p.GetContainerServer(remote)
	if err != nil {
		return err
	}

	_, _, err = server.GetContainer(containerName)
	if err != nil {
		return fmt.Errorf("unable to retrieve container %s: %s", containerName, err)
	}

	_, file, err := server.GetContainerFile(containerName, targetFile)
	if err != nil {
		return fmt.Errorf("unable to retrieve file %s:%s: %s", containerName, targetFile, err)
	}

	log.Printf("[DEBUG] Retrieved file: %#v", file)

	d.Set("container_name", containerName)
	d.Set("target_file", targetFile)
	d.Set("uid", file.UID)
	d.Set("gid", file.GID)
	d.Set("mode", file.Mode)

	return nil
}

func resourceLxdContainerFileDelete(d *schema.ResourceData, meta interface{}) error {
	p := meta.(*lxdProvider)
	v, targetFile := newFileIDFromResourceID(d.Id())
	remote, containerName, err := p.Config.ParseRemote(v)
	server, err := p.GetContainerServer(remote)
	if err != nil {
		return err
	}

	_, _, err = server.GetContainer(containerName)
	if err != nil {
		return fmt.Errorf("unable to retrieve container %s: %s", containerName, err)
	}

	err = containerDeleteFile(server, containerName, targetFile)
	if err != nil {
		return fmt.Errorf("unable to delete file %s:%s: %s", containerName, targetFile, err)
	}

	return nil
}

func resourceLxdContainerFileExists(d *schema.ResourceData, meta interface{}) (exists bool, err error) {
	p := meta.(*lxdProvider)
	v, targetFile := newFileIDFromResourceID(d.Id())
	remote, containerName, err := p.Config.ParseRemote(v)
	if err != nil {
		err = fmt.Errorf("unable to determine remote: %s", err)
		return
	}

	server, err := p.GetContainerServer(remote)
	if err != nil {
		return
	}

	_, _, err = server.GetContainer(containerName)
	if err != nil {
		err = fmt.Errorf("unable to retrieve container %s: %s", containerName, err)
		return
	}

	_, _, err = server.GetContainerFile(containerName, targetFile)
	if err != nil {
		// If the file could not be found, then it doesn't exist.
		// Ignore the error and return with exists set to false.
		if err.Error() == "not found" {
			err = nil
			return
		}

		return
	}

	exists = true

	return
}
<|endoftext|>


<commit_before>package pilosa

import (
	"fmt"
	"io"
	"sort"
	"time"

	"github.com/pkg/errors"
)

type recordImportManager struct {
	client *Client
}

func newRecordImportManager(client *Client) *recordImportManager {
	return &recordImportManager{
		client: client,
	}
}

type importWorkerChannels struct {
	records <-chan []Record
	errs    chan<- error
	status  chan<- ImportStatusUpdate
}

func (rim recordImportManager) run(field *Field, iterator RecordIterator, options ImportOptions) error {
	shardWidth := field.index.shardWidth
	threadCount := uint64(options.threadCount)
	recordChans := make([]chan []Record, threadCount)
	recordBufs := make([][]Record, threadCount)
	errChan := make(chan error)
	recordErrChan := make(chan error, 1)
	statusChan := options.statusChan

	if options.importRecordsFunction == nil {
		return errors.New("importRecords function is required")
	}

	for i := range recordChans {
		recordChans[i] = make(chan []Record, options.batchSize)
		recordBufs[i] = make([]Record, 0, 16)
		chans := importWorkerChannels{
			records: recordChans[i],
			errs:    errChan,
			status:  statusChan,
		}
		go recordImportWorker(i, rim.client, field, chans, &options)
	}

	var importErr error
	done := uint64(0)

	go func(it RecordIterator) {
		var record Record
		var err error
		for {
			record, err = it.NextRecord()
			if err != nil {
				if err == io.EOF {
					err = nil
				}
				break
			}
			shard := record.Shard(shardWidth)
			idx := shard % threadCount
			recordBufs[idx] = append(recordBufs[idx], record)
			if len(recordBufs[idx]) == cap(recordBufs[idx]) {
				recordChans[idx] <- recordBufs[idx]
				recordBufs[idx] = make([]Record, 0, 16)
			}
		}
		// send any trailing data
		for idx, buf := range recordBufs {
			if len(buf) > 0 {
				recordChans[idx] <- buf
				recordBufs[idx] = nil
			}
		}
		recordErrChan <- err
	}(iterator)

sendRecords:
	for done < threadCount {
		select {
		case workerErr := <-errChan:
			done++
			if workerErr != nil {
				importErr = workerErr
				break sendRecords
			}
		case recordErr := <-recordErrChan:
			for _, q := range recordChans {
				close(q)
			}
			if recordErr != nil {
				importErr = recordErr
				break sendRecords
			}
		}
	}

	return importErr
}

func recordImportWorker(id int, client *Client, field *Field, chans importWorkerChannels, options *ImportOptions) {
	var err error
	batchForShard := map[uint64][]Record{}
	statusChan := chans.status
	recordChan := chans.records
	errChan := chans.errs
	shardNodes := map[uint64][]fragmentNode{}

	defer func() {
		if r := recover(); r != nil {
			if err == nil {
				err = fmt.Errorf("worker %d panic: %v", id, r)
			}
		}
		errChan <- err
	}()

	state := &importState{}
	recordCount := 0
	batchSize := options.batchSize
	shardWidth := field.index.shardWidth

readRecords:
	for recordBatch := range recordChan {
		for _, record := range recordBatch {
			recordCount++
			shard := record.Shard(shardWidth)
			if batchForShard[shard] == nil {
				batchForShard[shard] = make([]Record, 0, batchSize)
			}
			batchForShard[shard] = append(batchForShard[shard], record)

			if recordCount >= batchSize {
				for shard, records := range batchForShard {
					if len(records) == 0 {
						continue
					}
					err = importRecords(id, client, field, shardNodes, shard, records, options, statusChan, state)
					if err != nil {
						break readRecords
					}
					batchForShard[shard] = batchForShard[shard][:0]
				}
				recordCount = 0
			}
		}
	}

	if err != nil {
		return
	}

	// import remaining records
	for shard, records := range batchForShard {
		if len(records) == 0 {
			continue
		}
		err = importRecords(id, client, field, shardNodes, shard, records, options, statusChan, state)
		if err != nil {
			break
		}
	}
}

func importRecords(id int, client *Client, field *Field,
	shardNodes map[uint64][]fragmentNode,
	shard uint64,
	records []Record,
	options *ImportOptions,
	statusChan chan<- ImportStatusUpdate,
	state *importState) error {

	var nodes []fragmentNode
	var ok bool
	var err error
	importFun := options.importRecordsFunction
	if nodes, ok = shardNodes[shard]; !ok {
		// if the data has row or column keys, send the data only to the coordinator
		if field.index.options.keys || field.options.keys {
			node, err := client.fetchCoordinatorNode()
			if err != nil {
				return err
			}
			nodes = []fragmentNode{node}
		} else {
			nodes, err = client.fetchFragmentNodes(field.index.name, shard)
			if err != nil {
				return err
			}
		}
	}
	tic := time.Now()
	if !options.skipSort {
		sort.Sort(recordSort(records))
	}
	err = importFun(field, shard, records, nodes, options, state)
	if err != nil {
		return err
	}
	took := time.Since(tic)
	if statusChan != nil {
		statusChan <- ImportStatusUpdate{
			ThreadID:      id,
			Shard:         shard,
			ImportedCount: len(records),
			Time:          took,
		}
	}
	return nil
}

// ImportStatusUpdate contains the import progress information.
type ImportStatusUpdate struct {
	ThreadID      int
	Shard         uint64
	ImportedCount int
	Time          time.Duration
}

type recordSort []Record

func (rc recordSort) Len() int {
	return len(rc)
}

func (rc recordSort) Swap(i, j int) {
	rc[i], rc[j] = rc[j], rc[i]
}

func (rc recordSort) Less(i, j int) bool {
	return rc[i].Less(rc[j])
}
<commit_msg>Don't use batch size for channel buffers<commit_after>package pilosa

import (
	"fmt"
	"io"
	"sort"
	"time"

	"github.com/pkg/errors"
)

type recordImportManager struct {
	client *Client
}

func newRecordImportManager(client *Client) *recordImportManager {
	return &recordImportManager{
		client: client,
	}
}

type importWorkerChannels struct {
	records <-chan []Record
	errs    chan<- error
	status  chan<- ImportStatusUpdate
}

func (rim recordImportManager) run(field *Field, iterator RecordIterator, options ImportOptions) error {
	shardWidth := field.index.shardWidth
	threadCount := uint64(options.threadCount)
	recordChans := make([]chan []Record, threadCount)
	recordBufs := make([][]Record, threadCount)
	errChan := make(chan error)
	recordErrChan := make(chan error, 1)
	statusChan := options.statusChan

	if options.importRecordsFunction == nil {
		return errors.New("importRecords function is required")
	}

	for i := range recordChans {
		recordChans[i] = make(chan []Record, 16)
		recordBufs[i] = make([]Record, 0, 16)
		chans := importWorkerChannels{
			records: recordChans[i],
			errs:    errChan,
			status:  statusChan,
		}
		go recordImportWorker(i, rim.client, field, chans, &options)
	}

	var importErr error
	done := uint64(0)

	go func(it RecordIterator) {
		var record Record
		var err error
		for {
			record, err = it.NextRecord()
			if err != nil {
				if err == io.EOF {
					err = nil
				}
				break
			}
			shard := record.Shard(shardWidth)
			idx := shard % threadCount
			recordBufs[idx] = append(recordBufs[idx], record)
			if len(recordBufs[idx]) == cap(recordBufs[idx]) {
				recordChans[idx] <- recordBufs[idx]
				recordBufs[idx] = make([]Record, 0, 16)
			}
		}
		// send any trailing data
		for idx, buf := range recordBufs {
			if len(buf) > 0 {
				recordChans[idx] <- buf
				recordBufs[idx] = nil
			}
		}
		recordErrChan <- err
	}(iterator)

sendRecords:
	for done < threadCount {
		select {
		case workerErr := <-errChan:
			done++
			if workerErr != nil {
				importErr = workerErr
				break sendRecords
			}
		case recordErr := <-recordErrChan:
			for _, q := range recordChans {
				close(q)
			}
			if recordErr != nil {
				importErr = recordErr
				break sendRecords
			}
		}
	}

	return importErr
}

func recordImportWorker(id int, client *Client, field *Field, chans importWorkerChannels, options *ImportOptions) {
	var err error
	batchForShard := map[uint64][]Record{}
	statusChan := chans.status
	recordChan := chans.records
	errChan := chans.errs
	shardNodes := map[uint64][]fragmentNode{}

	defer func() {
		if r := recover(); r != nil {
			if err == nil {
				err = fmt.Errorf("worker %d panic: %v", id, r)
			}
		}
		errChan <- err
	}()

	state := &importState{}
	recordCount := 0
	batchSize := options.batchSize
	shardWidth := field.index.shardWidth

readRecords:
	for recordBatch := range recordChan {
		for _, record := range recordBatch {
			recordCount++
			shard := record.Shard(shardWidth)
			if batchForShard[shard] == nil {
				batchForShard[shard] = make([]Record, 0, batchSize)
			}
			batchForShard[shard] = append(batchForShard[shard], record)

			if recordCount >= batchSize {
				for shard, records := range batchForShard {
					if len(records) == 0 {
						continue
					}
					err = importRecords(id, client, field, shardNodes, shard, records, options, statusChan, state)
					if err != nil {
						break readRecords
					}
					batchForShard[shard] = batchForShard[shard][:0]
				}
				recordCount = 0
			}
		}
	}

	if err != nil {
		return
	}

	// import remaining records
	for shard, records := range batchForShard {
		if len(records) == 0 {
			continue
		}
		err = importRecords(id, client, field, shardNodes, shard, records, options, statusChan, state)
		if err != nil {
			break
		}
	}
}

func importRecords(id int, client *Client, field *Field,
	shardNodes map[uint64][]fragmentNode,
	shard uint64,
	records []Record,
	options *ImportOptions,
	statusChan chan<- ImportStatusUpdate,
	state *importState) error {

	var nodes []fragmentNode
	var ok bool
	var err error
	importFun := options.importRecordsFunction
	if nodes, ok = shardNodes[shard]; !ok {
		// if the data has row or column keys, send the data only to the coordinator
		if field.index.options.keys || field.options.keys {
			node, err := client.fetchCoordinatorNode()
			if err != nil {
				return err
			}
			nodes = []fragmentNode{node}
		} else {
			nodes, err = client.fetchFragmentNodes(field.index.name, shard)
			if err != nil {
				return err
			}
		}
	}
	tic := time.Now()
	if !options.skipSort {
		sort.Sort(recordSort(records))
	}
	err = importFun(field, shard, records, nodes, options, state)
	if err != nil {
		return err
	}
	took := time.Since(tic)
	if statusChan != nil {
		statusChan <- ImportStatusUpdate{
			ThreadID:      id,
			Shard:         shard,
			ImportedCount: len(records),
			Time:          took,
		}
	}
	return nil
}

// ImportStatusUpdate contains the import progress information.
type ImportStatusUpdate struct {
	ThreadID      int
	Shard         uint64
	ImportedCount int
	Time          time.Duration
}

type recordSort []Record

func (rc recordSort) Len() int {
	return len(rc)
}

func (rc recordSort) Swap(i, j int) {
	rc[i], rc[j] = rc[j], rc[i]
}

func (rc recordSort) Less(i, j int) bool {
	return rc[i].Less(rc[j])
}
<|endoftext|>


<commit_before>package pilosa

import (
	"fmt"
	"io"
	"time"

	lru "github.com/hashicorp/golang-lru/simplelru"
	"github.com/pkg/errors"
)

type recordImportManager struct {
	client *Client
}

func newRecordImportManager(client *Client) *recordImportManager {
	return &recordImportManager{
		client: client,
	}
}

type importWorkerChannels struct {
	records <-chan Record
	errs    chan<- error
	status  chan<- ImportStatusUpdate
}

func (rim recordImportManager) run(field *Field, iterator RecordIterator, options ImportOptions) error {
	shardWidth := field.index.shardWidth
	threadCount := uint64(options.threadCount)
	recordChans := make([]chan Record, threadCount)
	errChan := make(chan error)
	recordErrChan := make(chan error, 1)
	statusChan := options.statusChan

	if options.importRecordsFunction == nil {
		return errors.New("importRecords function is required")
	}

	for i := range recordChans {
		recordChans[i] = make(chan Record, options.batchSize)
		chans := importWorkerChannels{
			records: recordChans[i],
			errs:    errChan,
			status:  statusChan,
		}
		go recordImportWorker(i, rim.client, field, chans, &options)
	}

	var importErr error
	done := uint64(0)

	go func(it RecordIterator) {
		var record Record
		var err error
		for {
			record, err = it.NextRecord()
			if err != nil {
				if err == io.EOF {
					err = nil
				}
				break
			}
			shard := record.Shard(shardWidth)
			recordChans[shard%threadCount] <- record
		}
		recordErrChan <- err
	}(iterator)

sendRecords:
	for done < threadCount {
		select {
		case workerErr := <-errChan:
			done++
			if workerErr != nil {
				importErr = workerErr
				break sendRecords
			}
		case recordErr := <-recordErrChan:
			for _, q := range recordChans {
				close(q)
			}
			if recordErr != nil {
				importErr = recordErr
				break sendRecords
			}
		}
	}

	return importErr
}

func recordImportWorker(id int, client *Client, field *Field, chans importWorkerChannels, options *ImportOptions) {
	var err error
	batchForShard := map[uint64][]Record{}
	statusChan := chans.status
	recordChan := chans.records
	errChan := chans.errs
	shardNodes := map[uint64][]fragmentNode{}

	defer func() {
		if r := recover(); r != nil {
			if err == nil {
				err = fmt.Errorf("worker %d panic: %v", id, r)
			}
		}
		errChan <- err
	}()

	rowKeyIDMap, err := lru.NewLRU(options.rowKeyCacheSize, nil)
	if err != nil {
		panic(errors.Wrap(err, "while creating rowKey to ID map"))
	}
	columnKeyIDMap, err := lru.NewLRU(options.columnKeyCacheSize, nil)
	if err != nil {
		panic(errors.Wrap(err, "while creating columnKey to ID map"))
	}
	state := &importState{
		rowKeyIDMap:    rowKeyIDMap,
		columnKeyIDMap: columnKeyIDMap,
	}

	recordCount := 0
	batchSize := options.batchSize
	shardWidth := field.index.shardWidth

readRecords:
	for record := range recordChan {
		recordCount++
		shard := record.Shard(shardWidth)
		batchForShard[shard] = append(batchForShard[shard], record)

		if recordCount >= batchSize {
			for shard, records := range batchForShard {
				if len(records) == 0 {
					continue
				}
				err = importRecords(id, client, field, shardNodes, shard, records, options, statusChan, state)
				if err != nil {
					break readRecords
				}
				batchForShard[shard] = nil
			}
			recordCount = 0
		}
	}

	if err != nil {
		return
	}

	// import remaining records
	for shard, records := range batchForShard {
		if len(records) == 0 {
			continue
		}
		err = importRecords(id, client, field, shardNodes, shard, records, options, statusChan, state)
		if err != nil {
			break
		}
	}
}

func importRecords(id int, client *Client, field *Field,
	shardNodes map[uint64][]fragmentNode,
	shard uint64,
	records []Record,
	options *ImportOptions,
	statusChan chan<- ImportStatusUpdate,
	state *importState) error {

	var nodes []fragmentNode
	var ok bool
	var err error
	importFun := options.importRecordsFunction
	if nodes, ok = shardNodes[shard]; !ok {
		// if the data has row or column keys, send the data only to the coordinator
		if field.index.options.keys || field.options.keys {
			node, err := client.fetchCoordinatorNode()
			if err != nil {
				return err
			}
			nodes = []fragmentNode{node}
		} else {
			nodes, err = client.fetchFragmentNodes(field.index.name, shard)
			if err != nil {
				return err
			}
		}
	}
	tic := time.Now()
	err = importFun(field, shard, records, nodes, options, state)
	if err != nil {
		return err
	}
	took := time.Since(tic)
	if statusChan != nil {
		statusChan <- ImportStatusUpdate{
			ThreadID:      id,
			Shard:         shard,
			ImportedCount: len(records),
			Time:          took,
		}
	}
	return nil
}

// ImportStatusUpdate contains the import progress information.
type ImportStatusUpdate struct {
	ThreadID      int
	Shard         uint64
	ImportedCount int
	Time          time.Duration
}

type recordSort []Record

func (rc recordSort) Len() int {
	return len(rc)
}

func (rc recordSort) Swap(i, j int) {
	rc[i], rc[j] = rc[j], rc[i]
}

func (rc recordSort) Less(i, j int) bool {
	return rc[i].Less(rc[j])
}
<commit_msg>reuse record buffers<commit_after>package pilosa

import (
	"fmt"
	"io"
	"time"

	lru "github.com/hashicorp/golang-lru/simplelru"
	"github.com/pkg/errors"
)

type recordImportManager struct {
	client *Client
}

func newRecordImportManager(client *Client) *recordImportManager {
	return &recordImportManager{
		client: client,
	}
}

type importWorkerChannels struct {
	records <-chan Record
	errs    chan<- error
	status  chan<- ImportStatusUpdate
}

func (rim recordImportManager) run(field *Field, iterator RecordIterator, options ImportOptions) error {
	shardWidth := field.index.shardWidth
	threadCount := uint64(options.threadCount)
	recordChans := make([]chan Record, threadCount)
	errChan := make(chan error)
	recordErrChan := make(chan error, 1)
	statusChan := options.statusChan

	if options.importRecordsFunction == nil {
		return errors.New("importRecords function is required")
	}

	for i := range recordChans {
		recordChans[i] = make(chan Record, options.batchSize)
		chans := importWorkerChannels{
			records: recordChans[i],
			errs:    errChan,
			status:  statusChan,
		}
		go recordImportWorker(i, rim.client, field, chans, &options)
	}

	var importErr error
	done := uint64(0)

	go func(it RecordIterator) {
		var record Record
		var err error
		for {
			record, err = it.NextRecord()
			if err != nil {
				if err == io.EOF {
					err = nil
				}
				break
			}
			shard := record.Shard(shardWidth)
			recordChans[shard%threadCount] <- record
		}
		recordErrChan <- err
	}(iterator)

sendRecords:
	for done < threadCount {
		select {
		case workerErr := <-errChan:
			done++
			if workerErr != nil {
				importErr = workerErr
				break sendRecords
			}
		case recordErr := <-recordErrChan:
			for _, q := range recordChans {
				close(q)
			}
			if recordErr != nil {
				importErr = recordErr
				break sendRecords
			}
		}
	}

	return importErr
}

func recordImportWorker(id int, client *Client, field *Field, chans importWorkerChannels, options *ImportOptions) {
	var err error
	batchForShard := map[uint64][]Record{}
	statusChan := chans.status
	recordChan := chans.records
	errChan := chans.errs
	shardNodes := map[uint64][]fragmentNode{}

	defer func() {
		if r := recover(); r != nil {
			if err == nil {
				err = fmt.Errorf("worker %d panic: %v", id, r)
			}
		}
		errChan <- err
	}()

	rowKeyIDMap, err := lru.NewLRU(options.rowKeyCacheSize, nil)
	if err != nil {
		panic(errors.Wrap(err, "while creating rowKey to ID map"))
	}
	columnKeyIDMap, err := lru.NewLRU(options.columnKeyCacheSize, nil)
	if err != nil {
		panic(errors.Wrap(err, "while creating columnKey to ID map"))
	}
	state := &importState{
		rowKeyIDMap:    rowKeyIDMap,
		columnKeyIDMap: columnKeyIDMap,
	}

	recordCount := 0
	batchSize := options.batchSize
	shardWidth := field.index.shardWidth

readRecords:
	for record := range recordChan {
		recordCount++
		shard := record.Shard(shardWidth)
		if batchForShard[shard] == nil {
			batchForShard[shard] = make([]Record, 0, batchSize)
		}
		batchForShard[shard] = append(batchForShard[shard], record)

		if recordCount >= batchSize {
			for shard, records := range batchForShard {
				if len(records) == 0 {
					continue
				}
				err = importRecords(id, client, field, shardNodes, shard, records, options, statusChan, state)
				if err != nil {
					break readRecords
				}
				batchForShard[shard] = batchForShard[shard][:0]
			}
			recordCount = 0
		}
	}

	if err != nil {
		return
	}

	// import remaining records
	for shard, records := range batchForShard {
		if len(records) == 0 {
			continue
		}
		err = importRecords(id, client, field, shardNodes, shard, records, options, statusChan, state)
		if err != nil {
			break
		}
	}
}

func importRecords(id int, client *Client, field *Field,
	shardNodes map[uint64][]fragmentNode,
	shard uint64,
	records []Record,
	options *ImportOptions,
	statusChan chan<- ImportStatusUpdate,
	state *importState) error {

	var nodes []fragmentNode
	var ok bool
	var err error
	importFun := options.importRecordsFunction
	if nodes, ok = shardNodes[shard]; !ok {
		// if the data has row or column keys, send the data only to the coordinator
		if field.index.options.keys || field.options.keys {
			node, err := client.fetchCoordinatorNode()
			if err != nil {
				return err
			}
			nodes = []fragmentNode{node}
		} else {
			nodes, err = client.fetchFragmentNodes(field.index.name, shard)
			if err != nil {
				return err
			}
		}
	}
	tic := time.Now()
	err = importFun(field, shard, records, nodes, options, state)
	if err != nil {
		return err
	}
	took := time.Since(tic)
	if statusChan != nil {
		statusChan <- ImportStatusUpdate{
			ThreadID:      id,
			Shard:         shard,
			ImportedCount: len(records),
			Time:          took,
		}
	}
	return nil
}

// ImportStatusUpdate contains the import progress information.
type ImportStatusUpdate struct {
	ThreadID      int
	Shard         uint64
	ImportedCount int
	Time          time.Duration
}

type recordSort []Record

func (rc recordSort) Len() int {
	return len(rc)
}

func (rc recordSort) Swap(i, j int) {
	rc[i], rc[j] = rc[j], rc[i]
}

func (rc recordSort) Less(i, j int) bool {
	return rc[i].Less(rc[j])
}
<|endoftext|>


<commit_before>package pinyin

import (
	"reflect"
	"testing"
)

type pinyinFunc func(string, Args) [][]string

func suite(s string, d map[Args][][]string, f pinyinFunc) (t *testing.T) {
	for a := range d {
		value, _ := d[a]
		v := f(s, a)
		if !reflect.DeepEqual(v, value) {
			t.Errorf("Expected %s, got %s", value, v)
		}
	}
	return t
}

func TestPinyin(t *testing.T) {
	hans := "中国人"
	test_data := map[Args][][]string{
		Args{}: [][]string{
			[]string{"zhong"},
			[]string{"guo"},
			[]string{"ren"},
		},
		Args{Style: NORMAL}: [][]string{
			[]string{"zhong"},
			[]string{"guo"},
			[]string{"ren"},
		},
		Args{Style: TONE}: [][]string{
			[]string{"zhōng"},
			[]string{"guó"},
			[]string{"rén"},
		},
		Args{Style: TONE2}: [][]string{
			[]string{"zho1ng"},
			[]string{"guo2"},
			[]string{"re2n"},
		},
		Args{Style: INITIALS}: [][]string{
			[]string{"zh"},
			[]string{"g"},
			[]string{"r"},
		},
		Args{Style: FIRST_LETTER}: [][]string{
			[]string{"z"},
			[]string{"g"},
			[]string{"r"},
		},
		Args{Style: FINALS}: [][]string{
			[]string{"ong"},
			[]string{"uo"},
			[]string{"en"},
		},
		Args{Style: FINALS_TONE}: [][]string{
			[]string{"ōng"},
			[]string{"uó"},
			[]string{"én"},
		},
		Args{Style: FINALS_TONE2}: [][]string{
			[]string{"o1ng"},
			[]string{"uo2"},
			[]string{"e2n"},
		},
		Args{Heteronym: true}: [][]string{
			[]string{"zhong", "zhong"},
			[]string{"guo"},
			[]string{"ren"},
		},
	}

	suite(hans, test_data, Pinyin)
}

func TestLazyPinyin(t *testing.T) {
	s := "中国人"
	v := LazyPinyin(s, Args{})
	value := []string{"zhong", "guo", "ren"}
	if !reflect.DeepEqual(v, value) {
		t.Errorf("Expected %s, got %s", value, v)
	}
}

func TestSlug(t *testing.T) {
	s := "中国人"
	v := Slug(s, Args{})
	value := "zhong-guo-ren"
	if v != value {
		t.Errorf("Expected %s, got %s", value, v)
	}
}
<commit_msg>update tests<commit_after>package pinyin

import (
	"reflect"
	"testing"
)

type pinyinFunc func(string, Args) [][]string

func testPinyin(s string, d map[Args][][]string, f pinyinFunc) (t *testing.T) {
	for a := range d {
		value, _ := d[a]
		v := f(s, a)
		if !reflect.DeepEqual(v, value) {
			t.Errorf("Expected %s, got %s", value, v)
		}
	}
	return t
}

func TestPinyin(t *testing.T) {
	hans := "中国人"
	test_data := map[Args][][]string{
		Args{}: [][]string{
			[]string{"zhong"},
			[]string{"guo"},
			[]string{"ren"},
		},
		Args{Style: NORMAL}: [][]string{
			[]string{"zhong"},
			[]string{"guo"},
			[]string{"ren"},
		},
		Args{Style: TONE}: [][]string{
			[]string{"zhōng"},
			[]string{"guó"},
			[]string{"rén"},
		},
		Args{Style: TONE2}: [][]string{
			[]string{"zho1ng"},
			[]string{"guo2"},
			[]string{"re2n"},
		},
		Args{Style: INITIALS}: [][]string{
			[]string{"zh"},
			[]string{"g"},
			[]string{"r"},
		},
		Args{Style: FIRST_LETTER}: [][]string{
			[]string{"z"},
			[]string{"g"},
			[]string{"r"},
		},
		Args{Style: FINALS}: [][]string{
			[]string{"ong"},
			[]string{"uo"},
			[]string{"en"},
		},
		Args{Style: FINALS_TONE}: [][]string{
			[]string{"ōng"},
			[]string{"uó"},
			[]string{"én"},
		},
		Args{Style: FINALS_TONE2}: [][]string{
			[]string{"o1ng"},
			[]string{"uo2"},
			[]string{"e2n"},
		},
		Args{Heteronym: true}: [][]string{
			[]string{"zhong", "zhong"},
			[]string{"guo"},
			[]string{"ren"},
		},
	}

	testPinyin(hans, test_data, Pinyin)
}

func TestLazyPinyin(t *testing.T) {
	s := "中国人"
	v := LazyPinyin(s, Args{})
	value := []string{"zhong", "guo", "ren"}
	if !reflect.DeepEqual(v, value) {
		t.Errorf("Expected %s, got %s", value, v)
	}
}

func TestSlug(t *testing.T) {
	s := "中国人"
	v := Slug(s, Args{})
	value := "zhong-guo-ren"
	if v != value {
		t.Errorf("Expected %s, got %s", value, v)
	}

	v = Slug(s, Args{Separator: ","})
	value = "zhong,guo,ren"
	if v != value {
		t.Errorf("Expected %s, got %s", value, v)
	}
}

func TestFinal(t *testing.T) {
	value := "an"
	v := final("an")
	if v != value {
		t.Errorf("Expected %s, got %s", value, v)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2016, Google Inc. All rights reserved.
//
// 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 rotation_support

import (
	"container/list"
	"fmt"
	"testing"
	"time"

	"github.com/jlmucb/cloudproxy/go/support_libraries/protected_objects"
)

func TestAddKeyEpoch(t *testing.T) {

	obj_type := "file"
	status := "active"
	nb := time.Now()
	validFor := 365 * 24 * time.Hour
	na := nb.Add(validFor)

	obj_1, err := protected_objects.CreateObject("/jlm/file/file1", 1,
		&obj_type, &status, &nb, &na, nil)
	if err != nil {
		t.Fatal("Can't create object")
	}
	fmt.Printf("Obj: %s\n", *obj_1.NotBefore)
	obj_type = "key"
	obj_2, _ := protected_objects.CreateObject("/jlm/key/key1", 1,
		&obj_type, &status, &nb, &na, nil)
	obj_3, _ := protected_objects.CreateObject("/jlm/key/key2", 1,
		&obj_type, &status, &nb, &na, nil)

	// add them to object list
	obj_list := list.New()
	err = protected_objects.AddObject(obj_list, *obj_1)
	if err != nil {
		t.Fatal("Can't add object")
	}
	_ = protected_objects.AddObject(obj_list, *obj_2)
	_ = protected_objects.AddObject(obj_list, *obj_3)

	newkey := []byte{0xff, 0xfe, 0xff, 0xfe, 0xff, 0xfe, 0xff, 0xfe,
		0x01, 0x02, 0x01, 0x02, 0x01, 0x02, 0x01, 0x02,
		0x07, 0x08, 0x07, 0x08, 0x07, 0x08, 0x07, 0x08,
		0xa6, 0xa5, 0xa6, 0xa5, 0xa6, 0xa5, 0xa6, 0xa5}

	oldkeyobj, newkeyobj, err := AddNewKeyEpoch(obj_list, "/jlm/key/key1",
		"key", "active", "active", nb.String(), na.String(), newkey)
	if err != nil {
		t.Fatal("Can't add new key epoch")
	}
	fmt.Printf("\n\n")
	if oldkeyobj == nil {
		fmt.Printf("No old key object\n")
	} else {
		fmt.Printf("Old key object:\n")
		protected_objects.PrintObject(oldkeyobj)
	}
	fmt.Printf("\n\n")
	if newkeyobj == nil {
		t.Fatal("Can't new key object is nil")
	}
	fmt.Printf("New key object:\n")
	protected_objects.PrintObject(newkeyobj)
	fmt.Printf("\n")

	oldkeyobj, newkeyobj, err = AddNewKeyEpoch(obj_list, "/jlm/key/key4",
		"key", "active", "active", nb.String(), na.String(), newkey)
	if err != nil {
		t.Fatal("Can't add new key epoch")
	}
	fmt.Printf("\n\n")
	if oldkeyobj == nil {
		fmt.Printf("No old key object\n")
	} else {
		fmt.Printf("Old key object:\n")
		protected_objects.PrintObject(oldkeyobj)
	}
	fmt.Printf("\n\n")
	if newkeyobj == nil {
		t.Fatal("Can't new key object is nil")
	}
	fmt.Printf("New key object:\n")
	protected_objects.PrintObject(newkeyobj)
	fmt.Printf("\n")
}

func TestAddAndRotate(t *testing.T) {

	obj_type := "file"
	status := "active"
	nb := time.Now()
	validFor := 365 * 24 * time.Hour
	na := nb.Add(validFor)

	protectorKeys := []byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0xa, 0xb, 0xc, 0xd, 0xe, 0xf,
		0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0xa, 0xb, 0xc, 0xd, 0xe, 0xf}

	obj_1, err := protected_objects.CreateObject("/jlm/file/file1", 1,
		&obj_type, &status, &nb, &na, nil)
	if err != nil {
		t.Fatal("Can't create object")
	}
	fmt.Printf("Obj: %s\n", *obj_1.NotBefore)
	obj_type = "key"
	obj_2, _ := protected_objects.CreateObject("/jlm/key/key1", 1,
		&obj_type, &status, &nb, &na, protectorKeys)
	obj_3, _ := protected_objects.CreateObject("/jlm/key/key2", 1,
		&obj_type, &status, &nb, &na, protectorKeys)

	// add them to object list
	obj_list := list.New()
	err = protected_objects.AddObject(obj_list, *obj_1)
	if err != nil {
		t.Fatal("Can't add object")
	}
	_ = protected_objects.AddObject(obj_list, *obj_2)
	_ = protected_objects.AddObject(obj_list, *obj_3)

	newkey := []byte{0xff, 0xfe, 0xff, 0xfe, 0xff, 0xfe, 0xff, 0xfe,
		0x01, 0x02, 0x01, 0x02, 0x01, 0x02, 0x01, 0x02,
		0x07, 0x08, 0x07, 0x08, 0x07, 0x08, 0x07, 0x08,
		0xa6, 0xa5, 0xa6, 0xa5, 0xa6, 0xa5, 0xa6, 0xa5}

	p_obj_1, err := protected_objects.MakeProtectedObject(*obj_1, "/jlm/key/key1", 1, protectorKeys)
	if err != nil {
		t.Fatal("Can't make protected object")
	}
	if p_obj_1 == nil {
		t.Fatal("Bad protected object")
	}

	p_obj_2, err := protected_objects.MakeProtectedObject(*obj_2, "/jlm/key/key2", 1, protectorKeys)
	if err != nil {
		t.Fatal("Can't make protected object")
	}
	if p_obj_2 == nil {
		t.Fatal("Bad protected object")
	}

	protected_obj_list := list.New()
	err = protected_objects.AddProtectedObject(protected_obj_list, *p_obj_1)
	if err != nil {
		t.Fatal("Can't add protected object")
	}
	err = protected_objects.AddProtectedObject(protected_obj_list, *p_obj_2)
	if err != nil {
		t.Fatal("Can't add protected object")
	}

	fmt.Printf("\n\n")
	fmt.Printf("Initial Objects\n")
	for e := obj_list.Front(); e != nil; e = e.Next() {
		o := e.Value.(protected_objects.ObjectMessage)
		protected_objects.PrintObject(&o)
	}
	fmt.Printf("\n\nInitial protected objects\n")
	for e := protected_obj_list.Front(); e != nil; e = e.Next() {
		o := e.Value.(protected_objects.ProtectedObjectMessage)
		protected_objects.PrintProtectedObject(&o)
	}
	fmt.Printf("\n\n")

	new_obj, err := AddAndRotateNewKeyEpoch("/jlm/key/key2", "key", "active", "active",
		nb.String(), na.String(), newkey, obj_list, protected_obj_list)
	if err != nil {
		fmt.Printf("Err: %s\n", err)
		t.Fatal("Can't AddAndRotateNewKeyEpoch")
	}
	fmt.Printf("\n\n")
	fmt.Printf("New key: %s, %d\n", new_obj.ObjId.ObjName, new_obj.ObjId.ObjEpoch)
	fmt.Printf("\n\n")
	fmt.Printf("Protected objects\n")
	for e := protected_obj_list.Front(); e != nil; e = e.Next() {
		o := e.Value.(protected_objects.ProtectedObjectMessage)
		protected_objects.PrintProtectedObject(&o)
	}
	fmt.Printf("\n\n")
	fmt.Printf("Objects\n")
	for e := obj_list.Front(); e != nil; e = e.Next() {
		o := e.Value.(protected_objects.ObjectMessage)
		protected_objects.PrintObject(&o)
	}
	fmt.Printf("\n\n")
	// Check we can open protected object with new protector
	protected_kids := protected_objects.FindProtectedObjects(protected_obj_list, *new_obj.ObjId.ObjName,
		*new_obj.ObjId.ObjEpoch)
	if err != nil {
		t.Fatal("Can't FindProtected kids")
	}
	e := protected_kids.Front()
	o := e.Value.(protected_objects.ProtectedObjectMessage)
	obj, err := protected_objects.RecoverProtectedObject(&o, new_obj.ObjVal)
	if err != nil || obj == nil {
		t.Fatal("Can't recover first kid")
	}
	fmt.Printf("\n\nRecovered:\n")
	protected_objects.PrintObject(obj)
}
<commit_msg>fix key size<commit_after>// Copyright (c) 2016, Google Inc. All rights reserved.
//
// 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 rotation_support

import (
	"container/list"
	"fmt"
	"testing"
	"time"

	"github.com/jlmucb/cloudproxy/go/support_libraries/protected_objects"
)

func TestAddKeyEpoch(t *testing.T) {

	obj_type := "file"
	status := "active"
	nb := time.Now()
	validFor := 365 * 24 * time.Hour
	na := nb.Add(validFor)

	obj_1, err := protected_objects.CreateObject("/jlm/file/file1", 1,
		&obj_type, &status, &nb, &na, nil)
	if err != nil {
		t.Fatal("Can't create object")
	}
	fmt.Printf("Obj: %s\n", *obj_1.NotBefore)
	obj_type = "key"
	obj_2, _ := protected_objects.CreateObject("/jlm/key/key1", 1,
		&obj_type, &status, &nb, &na, nil)
	obj_3, _ := protected_objects.CreateObject("/jlm/key/key2", 1,
		&obj_type, &status, &nb, &na, nil)

	// add them to object list
	obj_list := list.New()
	err = protected_objects.AddObject(obj_list, *obj_1)
	if err != nil {
		t.Fatal("Can't add object")
	}
	_ = protected_objects.AddObject(obj_list, *obj_2)
	_ = protected_objects.AddObject(obj_list, *obj_3)

	newkey := []byte{
		0xff, 0xfe, 0xff, 0xfe, 0xff, 0xfe, 0xff, 0xfe,
		0x01, 0x02, 0x01, 0x02, 0x01, 0x02, 0x01, 0x02,
		0x07, 0x08, 0x07, 0x08, 0x07, 0x08, 0x07, 0x08,
		0xa6, 0xa5, 0xa6, 0xa5, 0xa6, 0xa5, 0xa6, 0xa5}

	oldkeyobj, newkeyobj, err := AddNewKeyEpoch(obj_list, "/jlm/key/key1",
		"key", "active", "active", nb.String(), na.String(), newkey)
	if err != nil {
		t.Fatal("Can't add new key epoch")
	}
	fmt.Printf("\n\n")
	if oldkeyobj == nil {
		fmt.Printf("No old key object\n")
	} else {
		fmt.Printf("Old key object:\n")
		protected_objects.PrintObject(oldkeyobj)
	}
	fmt.Printf("\n\n")
	if newkeyobj == nil {
		t.Fatal("Can't new key object is nil")
	}
	fmt.Printf("New key object:\n")
	protected_objects.PrintObject(newkeyobj)
	fmt.Printf("\n")

	oldkeyobj, newkeyobj, err = AddNewKeyEpoch(obj_list, "/jlm/key/key4",
		"key", "active", "active", nb.String(), na.String(), newkey)
	if err != nil {
		t.Fatal("Can't add new key epoch")
	}
	fmt.Printf("\n\n")
	if oldkeyobj == nil {
		fmt.Printf("No old key object\n")
	} else {
		fmt.Printf("Old key object:\n")
		protected_objects.PrintObject(oldkeyobj)
	}
	fmt.Printf("\n\n")
	if newkeyobj == nil {
		t.Fatal("Can't new key object is nil")
	}
	fmt.Printf("New key object:\n")
	protected_objects.PrintObject(newkeyobj)
	fmt.Printf("\n")
}

func TestAddAndRotate(t *testing.T) {

	obj_type := "file"
	status := "active"
	nb := time.Now()
	validFor := 365 * 24 * time.Hour
	na := nb.Add(validFor)

	protectorKeys := []byte{
		0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0xa, 0xb, 0xc, 0xd, 0xe, 0xf,
		0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0xa, 0xb, 0xc, 0xd, 0xe, 0xf,
		0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0xa, 0xb, 0xc, 0xd, 0xe, 0xf,
	}

	obj_1, err := protected_objects.CreateObject("/jlm/file/file1", 1,
		&obj_type, &status, &nb, &na, nil)
	if err != nil {
		t.Fatal("Can't create object")
	}
	fmt.Printf("Obj: %s\n", *obj_1.NotBefore)
	obj_type = "key"
	obj_2, _ := protected_objects.CreateObject("/jlm/key/key1", 1,
		&obj_type, &status, &nb, &na, protectorKeys)
	obj_3, _ := protected_objects.CreateObject("/jlm/key/key2", 1,
		&obj_type, &status, &nb, &na, protectorKeys)

	// add them to object list
	obj_list := list.New()
	err = protected_objects.AddObject(obj_list, *obj_1)
	if err != nil {
		t.Fatal("Can't add object")
	}
	_ = protected_objects.AddObject(obj_list, *obj_2)
	_ = protected_objects.AddObject(obj_list, *obj_3)

	p_obj_1, err := protected_objects.MakeProtectedObject(*obj_1, "/jlm/key/key1", 1, protectorKeys)
	if err != nil {
		t.Fatal("Can't make protected object")
	}
	if p_obj_1 == nil {
		t.Fatal("Bad protected object")
	}

	p_obj_2, err := protected_objects.MakeProtectedObject(*obj_2, "/jlm/key/key2", 1, protectorKeys)
	if err != nil {
		t.Fatal("Can't make protected object")
	}
	if p_obj_2 == nil {
		t.Fatal("Bad protected object")
	}

	protected_obj_list := list.New()
	err = protected_objects.AddProtectedObject(protected_obj_list, *p_obj_1)
	if err != nil {
		t.Fatal("Can't add protected object")
	}
	err = protected_objects.AddProtectedObject(protected_obj_list, *p_obj_2)
	if err != nil {
		t.Fatal("Can't add protected object")
	}

	fmt.Printf("\n\n")
	fmt.Printf("Initial Objects\n")
	for e := obj_list.Front(); e != nil; e = e.Next() {
		o := e.Value.(protected_objects.ObjectMessage)
		protected_objects.PrintObject(&o)
	}
	fmt.Printf("\n\nInitial protected objects\n")
	for e := protected_obj_list.Front(); e != nil; e = e.Next() {
		o := e.Value.(protected_objects.ProtectedObjectMessage)
		protected_objects.PrintProtectedObject(&o)
	}
	fmt.Printf("\n\n")

	newkey := []byte{
		0xff, 0xfe, 0xff, 0xfe, 0xff, 0xfe, 0xff, 0xfe,
		0x01, 0x02, 0x01, 0x02, 0x01, 0x02, 0x01, 0x02,
		0x01, 0x02, 0x01, 0x02, 0x01, 0x02, 0x01, 0x02,
		0x07, 0x08, 0x07, 0x08, 0x07, 0x08, 0x07, 0x08,
		0x07, 0x08, 0x07, 0x08, 0x07, 0x08, 0x07, 0x08,
		0xa6, 0xa5, 0xa6, 0xa5, 0xa6, 0xa5, 0xa6, 0xa5}

	new_obj, err := AddAndRotateNewKeyEpoch("/jlm/key/key2", "key", "active", "active",
		nb.String(), na.String(), newkey, obj_list, protected_obj_list)
	if err != nil {
		fmt.Printf("Err: %s\n", err)
		t.Fatal("Can't AddAndRotateNewKeyEpoch")
	}
	fmt.Printf("\n\n")
	fmt.Printf("New key: %s, %d\n", new_obj.ObjId.ObjName, new_obj.ObjId.ObjEpoch)
	fmt.Printf("\n\n")
	fmt.Printf("Protected objects\n")
	for e := protected_obj_list.Front(); e != nil; e = e.Next() {
		o := e.Value.(protected_objects.ProtectedObjectMessage)
		protected_objects.PrintProtectedObject(&o)
	}
	fmt.Printf("\n\n")
	fmt.Printf("Objects\n")
	for e := obj_list.Front(); e != nil; e = e.Next() {
		o := e.Value.(protected_objects.ObjectMessage)
		protected_objects.PrintObject(&o)
	}
	fmt.Printf("\n\n")
	// Check we can open protected object with new protector
	protected_kids := protected_objects.FindProtectedObjects(protected_obj_list, *new_obj.ObjId.ObjName,
		*new_obj.ObjId.ObjEpoch)
	if err != nil {
		t.Fatal("Can't FindProtected kids")
	}
	e := protected_kids.Front()
	o := e.Value.(protected_objects.ProtectedObjectMessage)
	obj, err := protected_objects.RecoverProtectedObject(&o, new_obj.ObjVal)
	if err != nil || obj == nil {
		t.Fatal("Can't recover first kid")
	}
	fmt.Printf("\n\nRecovered:\n")
	protected_objects.PrintObject(obj)
}
<|endoftext|>


<commit_before>package gosql

import . "github.com/suboat/go-sql-kit"

type ValueFormatFunc func(string, string, interface{}) string

type SQLRule struct {
	RuleMapping
	ValueFormat ValueFormatFunc
}

func (s *SQLRule) SetValueFormat(f ValueFormatFunc) *SQLRule {
	if f != nil {
		s.ValueFormat = f
	}
	return s
}

func (s *SQLRule) ValueString(v *QueryValue) string {
	if v == nil {
	} else if s.ValueFormat == nil {
	} else if !s.IsAllowed(v.Field) {
	} else {
		return s.ValueFormat(v.Key, s.GetMapping(v.Field), v.Value)
	}
	return ""
}
<commit_msg>fix: query rule impl<commit_after>package gosql

import . "github.com/suboat/go-sql-kit"

type ValueFormatFunc func(string, string, interface{}) string

type SQLRule struct {
	RuleMapping
	ValueFormat ValueFormatFunc
}

func (s *SQLRule) SetValueFormat(f ValueFormatFunc) {
	s.ValueFormat = f
}

func (s *SQLRule) ValueString(v *QueryValue) string {
	if v == nil {
	} else if s.ValueFormat == nil {
	} else if !s.IsAllowed(v.Field) {
	} else {
		return s.ValueFormat(v.Key, s.GetMapping(v.Field), v.Value)
	}
	return ""
}
<|endoftext|>


<commit_before>package gotana

import (
	"bytes"
	"fmt"
	"github.com/PuerkitoBio/goquery"
	"golang.org/x/net/html"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	URL "net/url"
	"strings"
	"sync"
	"time"
)

const (
	EVENT_SCRAPER_OPENED     = "SCRAPER_OPENED"
	EVENT_SCRAPER_CLOSED     = "SCRAPER_CLOSED"
	EVENT_SAVEABLE_EXTRACTED = "SAVEABLE_EXTRACTED"
	TIMEOUT_DIALER           = time.Duration(time.Second * 30)
	TIMEOUT_REQUEST          = time.Duration(time.Second * 30)
	TIMEOUT_TLS              = time.Duration(time.Second * 10)
)

type SaveableItem interface {
	Scraper() *Scraper
	Validate() bool
	RecordData() []string
}

type ScraperMixin struct {
	Proxy ScrapedItem
}

func (s *ScraperMixin) SetProxy(proxy ScrapedItem) *ScraperMixin {
	s.Proxy = proxy
	return s
}

func (item ScraperMixin) Scraper() *Scraper {
	return item.Proxy.scraper
}

type recordWriter interface {
	Write(record []string) error
	Flush()
}

type ScrapingHandlerFunc func(ScrapedItem, chan<- SaveableItem)

func GetHref(t html.Token) (ok bool, href string) {
	for _, a := range t.Attr {
		if a.Key == "href" {
			href = a.Val
			ok = true
		}
	}

	return
}

type extensionParameters struct {
	scraper *Scraper
	item    SaveableItem
}

type Extractable interface {
	Extract(io.ReadCloser, func(string))
}

type LinkExtractor struct {
	Extractable
}

func (extractor *LinkExtractor) Extract(r io.ReadCloser, callback func(string)) {
	page := html.NewTokenizer(r)
	defer r.Close()

	for {
		tokenType := page.Next()
		if tokenType == html.ErrorToken {
			return
		}
		token := page.Token()
		if tokenType == html.StartTagToken && token.DataAtom.String() == "a" {
			ok, url := GetHref(token)
			if !ok {
				continue
			}
			callback(url)
		}
	}
}

type ScraperConfig struct {
	Project     string `required:"true"`
	TcpAddress  string
	OutFileName string
	Scrapers    []struct {
		RequestLimit int `required:"true"`
		Extractor    string
		Name         string `required:"true"`
		Url          string `required:"true"`
	}
}

type ScrapedItem struct {
	Url       string
	FinalUrl  string
	scraper   *Scraper
	BodyBytes []byte
}

func (proxy ScrapedItem) String() (result string) {
	result = fmt.Sprintf("Result of scraping: %s", proxy.Url)
	return
}

func (proxy ScrapedItem) CheckIfRedirected() bool {
	return proxy.Url != proxy.FinalUrl
}

func (proxy ScrapedItem) FinalResponseBody() (io.ReadCloser, error) {
	if proxy.CheckIfRedirected() {
		client := NewHTTPClient()
		response, err := client.Get(proxy.FinalUrl)
		if err != nil {
			return nil, err
		}
		bodyBytes, _ := ioutil.ReadAll(response.Body)
		proxy.BodyBytes = bodyBytes
	}
	return ioutil.NopCloser(bytes.NewBuffer(proxy.BodyBytes)), nil
}

func (proxy ScrapedItem) HTMLDocument() (document *goquery.Document, err error) {
	responseBody, err := proxy.FinalResponseBody()
	if err == nil {
		document, err = goquery.NewDocumentFromReader(responseBody)
	}

	return
}

type Scraper struct {
	crawled      int
	successful   int
	failed       int
	handler      ScrapingHandlerFunc
	fetchMutex   *sync.Mutex
	crawledMutex *sync.Mutex
	Name         string
	Domain       string
	BaseUrl      string
	CurrentUrl   string
	fetchedUrls  map[string]bool
	engine       *Engine
	extractor    Extractable
	chDone       chan struct{}
	chRequestUrl chan string
	requestLimit int
}

func (scraper *Scraper) MarkAsFetched(url string) {
	scraper.fetchMutex.Lock()
	defer scraper.fetchMutex.Unlock()

	scraper.CurrentUrl = url
	scraper.fetchedUrls[url] = true
}

func (scraper *Scraper) CheckIfShouldStop() (ok bool) {
	scraper.crawledMutex.Lock()
	defer scraper.crawledMutex.Unlock()
	stats := scraper.engine.Meta.ScraperStats[scraper.Name]

	if stats.crawled == scraper.engine.limitCrawl {
		Logger().Warningf("Crawl limit exceeded: %s", scraper)
		ok = true
	} else if stats.failed == scraper.engine.limitFail {
		Logger().Warningf("Fail limit exceeeded: %s", scraper)
		ok = true
	} else if stats.failed == 1 && scraper.crawled == 1 {
		Logger().Warningf("Base URL is corrupted: %s", scraper)
		ok = true
	}
	return
}

func (scraper *Scraper) CheckIfFetched(url string) (ok bool) {
	scraper.fetchMutex.Lock()
	defer scraper.fetchMutex.Unlock()

	_, ok = scraper.fetchedUrls[url]
	return
}

func (scraper *Scraper) CheckUrl(sourceUrl string) (ok bool, url string) {
	if strings.Contains(sourceUrl, scraper.Domain) && strings.Index(sourceUrl, "http") == 0 {
		url = sourceUrl
		ok = true
	} else if strings.Index(sourceUrl, "/") == 0 {
		url = scraper.BaseUrl + sourceUrl
		ok = true
	}
	return
}

func (scraper *Scraper) RunExtractor(resp *http.Response) {
	defer SilentRecover("EXTRACTOR")

	scraper.extractor.Extract(resp.Body, func(url string) {
		ok, url := scraper.CheckUrl(url)

		if ok {
			scraper.chRequestUrl <- url
		}
	})
}

func (scraper *Scraper) Stop() {
	Logger().Warningf("Stopping %s", scraper)
	scraper.engine.notifyExtensions(EVENT_SCRAPER_CLOSED,
		extensionParameters{scraper: scraper})

	scraper.chDone <- struct{}{}
	scraper.engine.wg.Done()
}

func (scraper *Scraper) Start() {
	scraper.engine.wg.Add(1)
	Logger().Infof("Starting: %s", scraper)
	scraper.engine.notifyExtensions(EVENT_SCRAPER_OPENED,
		extensionParameters{scraper: scraper})

	scraper.chRequestUrl <- scraper.BaseUrl
	duration := time.Duration(scraper.requestLimit)

	if scraper.requestLimit == 0 {
		duration = defaultRequestLimit()
	}

	limiter := time.Tick(time.Millisecond * duration)

	for {
		select {
		case url := <-scraper.chRequestUrl:
			<-limiter
			go scraper.Fetch(url)
		case <-scraper.chDone:
			Logger().Warningf("Stopped %s", scraper)
			scraper.engine.IncrFinishedCounter()
			scraper.engine.chDone <- struct{}{}
			return
		}
	}
	return
}

func (scraper *Scraper) Notify(url string, resp *http.Response) {
	scraper.engine.Meta.IncrScraped(scraper)
	scraper.engine.chScraped <- NewScrapedItem(url, scraper, resp)
}

func (scraper *Scraper) Fetch(url string) (resp *http.Response, err error) {
	if ok := scraper.CheckIfFetched(url); ok {
		return
	}
	scraper.MarkAsFetched(url)

	Logger().Infof("Fetching: %s", url)
	tic := time.Now()

	req, _ := http.NewRequest("GET", url, nil)
	req = scraper.engine.PrepareRequest(req)

	resp, err = NewHTTPClient().Do(req)

	statusCode := 0
	if err == nil {
		statusCode = resp.StatusCode
	}

	Logger().Debugf("[%d]Request to %s took: %s", statusCode, url, time.Since(tic))

	isSuccessful := (err == nil)

	scraper.engine.Meta.UpdateRequestStats(scraper, isSuccessful, req, resp)

	if err == nil {
		scraper.Notify(url, resp)
		scraper.RunExtractor(resp)
	} else {
		Logger().Warningf("Failed to crawl %s", url)
		Logger().Warning(err)
	}

	if scraper.CheckIfShouldStop() {
		scraper.Stop()
	}
	return
}

func (scraper *Scraper) SetHandler(handler ScrapingHandlerFunc) *Scraper {
	scraper.handler = handler
	return scraper
}

func (scraper *Scraper) String() (result string) {
	stats := scraper.engine.Meta.ScraperStats[scraper.Name]
	result = fmt.Sprintf("<Scraper: %s>. Crawled: %d, successful: %d, failed: %d. Items scraped: %d, saved: %d",
		scraper.Domain, stats.crawled, stats.successful, stats.failed, stats.scraped, stats.saved)
	return
}

func NewScraper(name string, sourceUrl string, requestLimit int, extractor Extractable) (s *Scraper) {
	parsed, err := URL.Parse(sourceUrl)
	if err != nil {
		Logger().Infof("Inappropriate URL: %s", sourceUrl)
		return
	}

	if extractor == nil {
		Logger().Warning("Switching to default extractor")
		extractor = defaultExtractor()
	}

	s = &Scraper{
		Name:         name,
		Domain:       parsed.Host,
		BaseUrl:      sourceUrl,
		fetchedUrls:  make(map[string]bool),
		crawledMutex: &sync.Mutex{},
		fetchMutex:   &sync.Mutex{},
		extractor:    extractor,
		chDone:       make(chan struct{}),
		chRequestUrl: make(chan string, 5),
		requestLimit: requestLimit,
	}
	return
}

func NewScrapedItem(url string, scraper *Scraper, resp *http.Response) ScrapedItem {
	bodyBytes, _ := ioutil.ReadAll(resp.Body)
	resp.Body = ioutil.NopCloser(bytes.NewBuffer(bodyBytes))

	return ScrapedItem{
		BodyBytes: bodyBytes,
		FinalUrl:  resp.Request.URL.String(),
		Url:       url,
		scraper:   scraper,
	}
}

func NewHTTPClient() (client *http.Client) {
	client = &http.Client{
		Timeout: TIMEOUT_REQUEST,
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout: TIMEOUT_DIALER,
			}).Dial,
			TLSHandshakeTimeout: TIMEOUT_TLS,
		},
	}
	return
}

func defaultExtractor() Extractable {
	return &LinkExtractor{}
}

func defaultRequestLimit() time.Duration {
	return time.Duration(1)
}

func SaveItem(item SaveableItem, writer recordWriter) {
	if writer == nil {
		return
	}

	if !item.Validate() {
		Logger().Warning("Item is not valid. Skipping...")
		return
	}

	defer writer.Flush()
	writer.Write(item.RecordData())
}

func NewSpiderConfig(file string) (config *ScraperConfig) {
	config = &ScraperConfig{}
	ProcessFile(config, file)
	return
}
<commit_msg>added trimhash<commit_after>package gotana

import (
	"bytes"
	"fmt"
	"github.com/PuerkitoBio/goquery"
	"golang.org/x/net/html"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	URL "net/url"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	EVENT_SCRAPER_OPENED     = "SCRAPER_OPENED"
	EVENT_SCRAPER_CLOSED     = "SCRAPER_CLOSED"
	EVENT_SAVEABLE_EXTRACTED = "SAVEABLE_EXTRACTED"
	TIMEOUT_DIALER           = time.Duration(time.Second * 30)
	TIMEOUT_REQUEST          = time.Duration(time.Second * 30)
	TIMEOUT_TLS              = time.Duration(time.Second * 10)
)

type SaveableItem interface {
	Scraper() *Scraper
	Validate() bool
	RecordData() []string
}

type ScraperMixin struct {
	Proxy ScrapedItem
}

func (s *ScraperMixin) SetProxy(proxy ScrapedItem) *ScraperMixin {
	s.Proxy = proxy
	return s
}

func (item ScraperMixin) Scraper() *Scraper {
	return item.Proxy.scraper
}

type recordWriter interface {
	Write(record []string) error
	Flush()
}

type ScrapingHandlerFunc func(ScrapedItem, chan<- SaveableItem)

func getHref(t html.Token) (ok bool, href string) {
	for _, a := range t.Attr {
		if a.Key == "href" {
			href = a.Val
			ok = true
		}
	}

	return
}

func trimHash(s string) string {
	if strings.Contains(s, "#") {
		var index int
		for i, str := range s {
			if strconv.QuoteRune(str) == "'#'" {
				index = i
				break
			}
		}
		return s[:index]
	}
	return s
}

type extensionParameters struct {
	scraper *Scraper
	item    SaveableItem
}

type Extractable interface {
	Extract(io.ReadCloser, func(string))
}

type LinkExtractor struct {
	Extractable
}

func (extractor *LinkExtractor) Extract(r io.ReadCloser, callback func(string)) {
	page := html.NewTokenizer(r)
	defer r.Close()

	for {
		tokenType := page.Next()
		if tokenType == html.ErrorToken {
			return
		}
		token := page.Token()
		if tokenType == html.StartTagToken && token.DataAtom.String() == "a" {
			ok, url := getHref(token)
			if !ok {
				continue
			}
			callback(trimHash(url))
		}
	}
}

type ScraperConfig struct {
	Project     string `required:"true"`
	TcpAddress  string
	OutFileName string
	Scrapers    []struct {
		RequestLimit int `required:"true"`
		Extractor    string
		Name         string `required:"true"`
		Url          string `required:"true"`
	}
}

type ScrapedItem struct {
	Url       string
	FinalUrl  string
	scraper   *Scraper
	BodyBytes []byte
}

func (proxy ScrapedItem) String() (result string) {
	result = fmt.Sprintf("Result of scraping: %s", proxy.Url)
	return
}

func (proxy ScrapedItem) CheckIfRedirected() bool {
	return proxy.Url != proxy.FinalUrl
}

func (proxy ScrapedItem) FinalResponseBody() (io.ReadCloser, error) {
	if proxy.CheckIfRedirected() {
		client := NewHTTPClient()
		response, err := client.Get(proxy.FinalUrl)
		if err != nil {
			return nil, err
		}
		bodyBytes, _ := ioutil.ReadAll(response.Body)
		proxy.BodyBytes = bodyBytes
	}
	return ioutil.NopCloser(bytes.NewBuffer(proxy.BodyBytes)), nil
}

func (proxy ScrapedItem) HTMLDocument() (document *goquery.Document, err error) {
	responseBody, err := proxy.FinalResponseBody()
	if err == nil {
		document, err = goquery.NewDocumentFromReader(responseBody)
	}

	return
}

type Scraper struct {
	crawled      int
	successful   int
	failed       int
	handler      ScrapingHandlerFunc
	fetchMutex   *sync.Mutex
	crawledMutex *sync.Mutex
	Name         string
	Domain       string
	BaseUrl      string
	CurrentUrl   string
	fetchedUrls  map[string]bool
	engine       *Engine
	extractor    Extractable
	chDone       chan struct{}
	chRequestUrl chan string
	requestLimit int
}

func (scraper *Scraper) MarkAsFetched(url string) {
	scraper.fetchMutex.Lock()
	defer scraper.fetchMutex.Unlock()

	scraper.CurrentUrl = url
	scraper.fetchedUrls[url] = true
}

func (scraper *Scraper) CheckIfShouldStop() (ok bool) {
	scraper.crawledMutex.Lock()
	defer scraper.crawledMutex.Unlock()
	stats := scraper.engine.Meta.ScraperStats[scraper.Name]

	if stats.crawled == scraper.engine.limitCrawl {
		Logger().Warningf("Crawl limit exceeded: %s", scraper)
		ok = true
	} else if stats.failed == scraper.engine.limitFail {
		Logger().Warningf("Fail limit exceeeded: %s", scraper)
		ok = true
	} else if stats.failed == 1 && scraper.crawled == 1 {
		Logger().Warningf("Base URL is corrupted: %s", scraper)
		ok = true
	}
	return
}

func (scraper *Scraper) CheckIfFetched(url string) (ok bool) {
	scraper.fetchMutex.Lock()
	defer scraper.fetchMutex.Unlock()

	_, ok = scraper.fetchedUrls[url]
	return
}

func (scraper *Scraper) CheckUrl(sourceUrl string) (ok bool, url string) {
	if strings.Contains(sourceUrl, scraper.Domain) && strings.Index(sourceUrl, "http") == 0 {
		url = sourceUrl
		ok = true
	} else if strings.Index(sourceUrl, "/") == 0 {
		url = scraper.BaseUrl + sourceUrl
		ok = true
	}
	return
}

func (scraper *Scraper) RunExtractor(resp *http.Response) {
	defer SilentRecover("EXTRACTOR")

	scraper.extractor.Extract(resp.Body, func(url string) {
		ok, url := scraper.CheckUrl(url)

		if ok {
			scraper.chRequestUrl <- url
		}
	})
}

func (scraper *Scraper) Stop() {
	Logger().Warningf("Stopping %s", scraper)
	scraper.engine.notifyExtensions(EVENT_SCRAPER_CLOSED,
		extensionParameters{scraper: scraper})

	scraper.chDone <- struct{}{}
	scraper.engine.wg.Done()
}

func (scraper *Scraper) Start() {
	scraper.engine.wg.Add(1)
	Logger().Infof("Starting: %s", scraper)
	scraper.engine.notifyExtensions(EVENT_SCRAPER_OPENED,
		extensionParameters{scraper: scraper})

	scraper.chRequestUrl <- scraper.BaseUrl
	duration := time.Duration(scraper.requestLimit)

	if scraper.requestLimit == 0 {
		duration = defaultRequestLimit()
	}

	limiter := time.Tick(time.Millisecond * duration)

	for {
		select {
		case url := <-scraper.chRequestUrl:
			<-limiter
			go scraper.Fetch(url)
		case <-scraper.chDone:
			Logger().Warningf("Stopped %s", scraper)
			scraper.engine.IncrFinishedCounter()
			scraper.engine.chDone <- struct{}{}
			return
		}
	}
	return
}

func (scraper *Scraper) Notify(url string, resp *http.Response) {
	scraper.engine.Meta.IncrScraped(scraper)
	scraper.engine.chScraped <- NewScrapedItem(url, scraper, resp)
}

func (scraper *Scraper) Fetch(url string) (resp *http.Response, err error) {
	if ok := scraper.CheckIfFetched(url); ok {
		return
	}
	scraper.MarkAsFetched(url)

	Logger().Infof("Fetching: %s", url)
	tic := time.Now()

	req, _ := http.NewRequest("GET", url, nil)
	req = scraper.engine.PrepareRequest(req)

	resp, err = NewHTTPClient().Do(req)

	statusCode := 0
	if err == nil {
		statusCode = resp.StatusCode
	}

	Logger().Debugf("[%d]Request to %s took: %s", statusCode, url, time.Since(tic))

	isSuccessful := (err == nil)

	scraper.engine.Meta.UpdateRequestStats(scraper, isSuccessful, req, resp)

	if err == nil {
		scraper.Notify(url, resp)
		scraper.RunExtractor(resp)
	} else {
		Logger().Warningf("Failed to crawl %s", url)
		Logger().Warning(err)
	}

	if scraper.CheckIfShouldStop() {
		scraper.Stop()
	}
	return
}

func (scraper *Scraper) SetHandler(handler ScrapingHandlerFunc) *Scraper {
	scraper.handler = handler
	return scraper
}

func (scraper *Scraper) String() (result string) {
	stats := scraper.engine.Meta.ScraperStats[scraper.Name]
	result = fmt.Sprintf("<Scraper: %s>. Crawled: %d, successful: %d, failed: %d. Items scraped: %d, saved: %d",
		scraper.Domain, stats.crawled, stats.successful, stats.failed, stats.scraped, stats.saved)
	return
}

func NewScraper(name string, sourceUrl string, requestLimit int, extractor Extractable) (s *Scraper) {
	parsed, err := URL.Parse(sourceUrl)
	if err != nil {
		Logger().Infof("Inappropriate URL: %s", sourceUrl)
		return
	}

	if extractor == nil {
		Logger().Warning("Switching to default extractor")
		extractor = defaultExtractor()
	}

	s = &Scraper{
		Name:         name,
		Domain:       parsed.Host,
		BaseUrl:      sourceUrl,
		fetchedUrls:  make(map[string]bool),
		crawledMutex: &sync.Mutex{},
		fetchMutex:   &sync.Mutex{},
		extractor:    extractor,
		chDone:       make(chan struct{}),
		chRequestUrl: make(chan string, 5),
		requestLimit: requestLimit,
	}
	return
}

func NewScrapedItem(url string, scraper *Scraper, resp *http.Response) ScrapedItem {
	bodyBytes, _ := ioutil.ReadAll(resp.Body)
	resp.Body = ioutil.NopCloser(bytes.NewBuffer(bodyBytes))

	return ScrapedItem{
		BodyBytes: bodyBytes,
		FinalUrl:  resp.Request.URL.String(),
		Url:       url,
		scraper:   scraper,
	}
}

func NewHTTPClient() (client *http.Client) {
	client = &http.Client{
		Timeout: TIMEOUT_REQUEST,
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout: TIMEOUT_DIALER,
			}).Dial,
			TLSHandshakeTimeout: TIMEOUT_TLS,
		},
	}
	return
}

func defaultExtractor() Extractable {
	return &LinkExtractor{}
}

func defaultRequestLimit() time.Duration {
	return time.Duration(1)
}

func SaveItem(item SaveableItem, writer recordWriter) {
	if writer == nil {
		return
	}

	if !item.Validate() {
		Logger().Warning("Item is not valid. Skipping...")
		return
	}

	defer writer.Flush()
	writer.Write(item.RecordData())
}

func NewSpiderConfig(file string) (config *ScraperConfig) {
	config = &ScraperConfig{}
	ProcessFile(config, file)
	return
}
<|endoftext|>


<commit_before>package main

import (
	"net/http"
	"testing"
)

func TestSynthesizeResponse(t *testing.T) {

	req, err := http.NewRequest("GET", "http://example.com", nil)
	expect(t, err, nil)

	sr := synthesizeResponse(req, "./examples/middleware/synthetic_service/synthetic.py")

	expect(t, sr.StatusCode, 200)
}
<commit_msg>test when middleware not provided<commit_after>package main

import (
	"net/http"
	"testing"
)

func TestSynthesizeResponse(t *testing.T) {

	req, err := http.NewRequest("GET", "http://example.com", nil)
	expect(t, err, nil)

	sr := synthesizeResponse(req, "./examples/middleware/synthetic_service/synthetic.py")

	expect(t, sr.StatusCode, 200)
}

func TestSynthesizeResponseWOMiddleware(t *testing.T) {

	req, err := http.NewRequest("GET", "http://example.com", nil)
	expect(t, err, nil)

	sr := synthesizeResponse(req, "")

	expect(t, sr.StatusCode, 428)
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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 deployment

import (
	"flag"
	"fmt"
	"testing"
	"time"

	fedv1 "k8s.io/kubernetes/federation/apis/federation/v1beta1"
	fake_fedclientset "k8s.io/kubernetes/federation/client/clientset_generated/federation_release_1_5/fake"
	. "k8s.io/kubernetes/federation/pkg/federation-controller/util/test"
	"k8s.io/kubernetes/pkg/api/meta"
	apiv1 "k8s.io/kubernetes/pkg/api/v1"
	extensionsv1 "k8s.io/kubernetes/pkg/apis/extensions/v1beta1"
	kubeclientset "k8s.io/kubernetes/pkg/client/clientset_generated/release_1_5"
	fake_kubeclientset "k8s.io/kubernetes/pkg/client/clientset_generated/release_1_5/fake"
	"k8s.io/kubernetes/pkg/runtime"

	"github.com/stretchr/testify/assert"
)

func TestParseFederationDeploymentPreference(t *testing.T) {
	successPrefs := []string{
		`{"rebalance": true,
		  "clusters": {
		    "k8s-1": {"minReplicas": 10, "maxReplicas": 20, "weight": 2},
		    "*": {"weight": 1}
		}}`,
	}
	failedPrefes := []string{
		`{`, // bad json
	}

	rs := newDeploymentWithReplicas("d-1", 100)
	accessor, _ := meta.Accessor(rs)
	anno := accessor.GetAnnotations()
	if anno == nil {
		anno = make(map[string]string)
		accessor.SetAnnotations(anno)
	}
	for _, prefString := range successPrefs {
		anno[FedDeploymentPreferencesAnnotation] = prefString
		pref, err := parseFederationDeploymentPreference(rs)
		assert.NotNil(t, pref)
		assert.Nil(t, err)
	}
	for _, prefString := range failedPrefes {
		anno[FedDeploymentPreferencesAnnotation] = prefString
		pref, err := parseFederationDeploymentPreference(rs)
		assert.Nil(t, pref)
		assert.NotNil(t, err)
	}
}

func TestDeploymentController(t *testing.T) {
	flag.Set("logtostderr", "true")
	flag.Set("v", "5")
	flag.Parse()

	deploymentReviewDelay = 500 * time.Millisecond
	clusterAvailableDelay = 100 * time.Millisecond
	clusterUnavailableDelay = 100 * time.Millisecond
	allDeploymentReviewDelay = 500 * time.Millisecond

	cluster1 := NewCluster("cluster1", apiv1.ConditionTrue)
	cluster2 := NewCluster("cluster2", apiv1.ConditionTrue)

	fakeClient := &fake_fedclientset.Clientset{}
	RegisterFakeList("clusters", &fakeClient.Fake, &fedv1.ClusterList{Items: []fedv1.Cluster{*cluster1}})
	deploymentsWatch := RegisterFakeWatch("deployments", &fakeClient.Fake)
	clusterWatch := RegisterFakeWatch("clusters", &fakeClient.Fake)

	cluster1Client := &fake_kubeclientset.Clientset{}
	cluster1Watch := RegisterFakeWatch("deployments", &cluster1Client.Fake)
	_ = RegisterFakeWatch("pods", &cluster1Client.Fake)
	RegisterFakeList("deployments", &cluster1Client.Fake, &extensionsv1.DeploymentList{Items: []extensionsv1.Deployment{}})
	cluster1CreateChan := RegisterFakeCopyOnCreate("deployments", &cluster1Client.Fake, cluster1Watch)
	cluster1UpdateChan := RegisterFakeCopyOnUpdate("deployments", &cluster1Client.Fake, cluster1Watch)

	cluster2Client := &fake_kubeclientset.Clientset{}
	cluster2Watch := RegisterFakeWatch("deployments", &cluster2Client.Fake)
	_ = RegisterFakeWatch("pods", &cluster2Client.Fake)
	RegisterFakeList("deployments", &cluster2Client.Fake, &extensionsv1.DeploymentList{Items: []extensionsv1.Deployment{}})
	cluster2CreateChan := RegisterFakeCopyOnCreate("deployments", &cluster2Client.Fake, cluster2Watch)

	deploymentController := NewDeploymentController(fakeClient)
	clientFactory := func(cluster *fedv1.Cluster) (kubeclientset.Interface, error) {
		switch cluster.Name {
		case cluster1.Name:
			return cluster1Client, nil
		case cluster2.Name:
			return cluster2Client, nil
		default:
			return nil, fmt.Errorf("Unknown cluster")
		}
	}
	ToFederatedInformerForTestOnly(deploymentController.fedDeploymentInformer).SetClientFactory(clientFactory)
	ToFederatedInformerForTestOnly(deploymentController.fedPodInformer).SetClientFactory(clientFactory)

	stop := make(chan struct{})
	go deploymentController.Run(5, stop)

	// Create deployment. Expect to see it in cluster1.
	dep1 := newDeploymentWithReplicas("depA", 6)
	deploymentsWatch.Add(dep1)
	checkDeployment := func(base *extensionsv1.Deployment, replicas int32) CheckingFunction {
		return func(obj runtime.Object) error {
			if obj == nil {
				return fmt.Errorf("Observed object is nil")
			}
			d := obj.(*extensionsv1.Deployment)
			if err := CompareObjectMeta(base.ObjectMeta, d.ObjectMeta); err != nil {
				return err
			}
			if replicas != *d.Spec.Replicas {
				return fmt.Errorf("Replica count is different expected:%d observed:%d", replicas, *d.Spec.Replicas)
			}
			return nil
		}
	}
	assert.NoError(t, CheckObjectFromChan(cluster1CreateChan, checkDeployment(dep1, *dep1.Spec.Replicas)))

	// Increase replica count. Expect to see the update in cluster1.
	newRep := int32(8)
	dep1.Spec.Replicas = &newRep
	deploymentsWatch.Modify(dep1)
	assert.NoError(t, CheckObjectFromChan(cluster1UpdateChan, checkDeployment(dep1, *dep1.Spec.Replicas)))

	// Add new cluster. Although rebalance = false, no pods have been created yet so it should
	// rebalance anyway.
	clusterWatch.Add(cluster2)
	assert.NoError(t, CheckObjectFromChan(cluster1UpdateChan, checkDeployment(dep1, *dep1.Spec.Replicas/2)))
	assert.NoError(t, CheckObjectFromChan(cluster2CreateChan, checkDeployment(dep1, *dep1.Spec.Replicas/2)))

	// Add new deployment with non-default replica placement preferences.
	dep2 := newDeploymentWithReplicas("deployment2", 9)
	dep2.Annotations = make(map[string]string)
	dep2.Annotations[FedDeploymentPreferencesAnnotation] = `{"rebalance": true,
		  "clusters": {
		    "cluster1": {"weight": 2},
		    "cluster2": {"weight": 1}
		}}`
	deploymentsWatch.Add(dep2)
	assert.NoError(t, CheckObjectFromChan(cluster1CreateChan, checkDeployment(dep2, 6)))
	assert.NoError(t, CheckObjectFromChan(cluster2CreateChan, checkDeployment(dep2, 3)))
}

func GetDeploymentFromChan(c chan runtime.Object) *extensionsv1.Deployment {
	secret := GetObjectFromChan(c).(*extensionsv1.Deployment)
	return secret
}

func newDeploymentWithReplicas(name string, replicas int32) *extensionsv1.Deployment {
	return &extensionsv1.Deployment{
		ObjectMeta: apiv1.ObjectMeta{
			Name:      name,
			Namespace: apiv1.NamespaceDefault,
			SelfLink:  "/api/v1/namespaces/default/deployments/name",
		},
		Spec: extensionsv1.DeploymentSpec{
			Replicas: &replicas,
		},
	}
}
<commit_msg>Wait for deployment in the storage in deployment controller test<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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 deployment

import (
	"flag"
	"fmt"
	"testing"
	"time"

	fedv1 "k8s.io/kubernetes/federation/apis/federation/v1beta1"
	fake_fedclientset "k8s.io/kubernetes/federation/client/clientset_generated/federation_release_1_5/fake"
	. "k8s.io/kubernetes/federation/pkg/federation-controller/util/test"
	"k8s.io/kubernetes/pkg/api/meta"
	apiv1 "k8s.io/kubernetes/pkg/api/v1"
	extensionsv1 "k8s.io/kubernetes/pkg/apis/extensions/v1beta1"
	kubeclientset "k8s.io/kubernetes/pkg/client/clientset_generated/release_1_5"
	fake_kubeclientset "k8s.io/kubernetes/pkg/client/clientset_generated/release_1_5/fake"
	"k8s.io/kubernetes/pkg/runtime"
	"k8s.io/kubernetes/pkg/types"
	"k8s.io/kubernetes/pkg/util/wait"

	"github.com/stretchr/testify/assert"
)

func TestParseFederationDeploymentPreference(t *testing.T) {
	successPrefs := []string{
		`{"rebalance": true,
		  "clusters": {
		    "k8s-1": {"minReplicas": 10, "maxReplicas": 20, "weight": 2},
		    "*": {"weight": 1}
		}}`,
	}
	failedPrefes := []string{
		`{`, // bad json
	}

	rs := newDeploymentWithReplicas("d-1", 100)
	accessor, _ := meta.Accessor(rs)
	anno := accessor.GetAnnotations()
	if anno == nil {
		anno = make(map[string]string)
		accessor.SetAnnotations(anno)
	}
	for _, prefString := range successPrefs {
		anno[FedDeploymentPreferencesAnnotation] = prefString
		pref, err := parseFederationDeploymentPreference(rs)
		assert.NotNil(t, pref)
		assert.Nil(t, err)
	}
	for _, prefString := range failedPrefes {
		anno[FedDeploymentPreferencesAnnotation] = prefString
		pref, err := parseFederationDeploymentPreference(rs)
		assert.Nil(t, pref)
		assert.NotNil(t, err)
	}
}

func TestDeploymentController(t *testing.T) {
	flag.Set("logtostderr", "true")
	flag.Set("v", "5")
	flag.Parse()

	deploymentReviewDelay = 500 * time.Millisecond
	clusterAvailableDelay = 100 * time.Millisecond
	clusterUnavailableDelay = 100 * time.Millisecond
	allDeploymentReviewDelay = 500 * time.Millisecond

	cluster1 := NewCluster("cluster1", apiv1.ConditionTrue)
	cluster2 := NewCluster("cluster2", apiv1.ConditionTrue)

	fakeClient := &fake_fedclientset.Clientset{}
	RegisterFakeList("clusters", &fakeClient.Fake, &fedv1.ClusterList{Items: []fedv1.Cluster{*cluster1}})
	deploymentsWatch := RegisterFakeWatch("deployments", &fakeClient.Fake)
	clusterWatch := RegisterFakeWatch("clusters", &fakeClient.Fake)

	cluster1Client := &fake_kubeclientset.Clientset{}
	cluster1Watch := RegisterFakeWatch("deployments", &cluster1Client.Fake)
	_ = RegisterFakeWatch("pods", &cluster1Client.Fake)
	RegisterFakeList("deployments", &cluster1Client.Fake, &extensionsv1.DeploymentList{Items: []extensionsv1.Deployment{}})
	cluster1CreateChan := RegisterFakeCopyOnCreate("deployments", &cluster1Client.Fake, cluster1Watch)
	cluster1UpdateChan := RegisterFakeCopyOnUpdate("deployments", &cluster1Client.Fake, cluster1Watch)

	cluster2Client := &fake_kubeclientset.Clientset{}
	cluster2Watch := RegisterFakeWatch("deployments", &cluster2Client.Fake)
	_ = RegisterFakeWatch("pods", &cluster2Client.Fake)
	RegisterFakeList("deployments", &cluster2Client.Fake, &extensionsv1.DeploymentList{Items: []extensionsv1.Deployment{}})
	cluster2CreateChan := RegisterFakeCopyOnCreate("deployments", &cluster2Client.Fake, cluster2Watch)

	deploymentController := NewDeploymentController(fakeClient)
	clientFactory := func(cluster *fedv1.Cluster) (kubeclientset.Interface, error) {
		switch cluster.Name {
		case cluster1.Name:
			return cluster1Client, nil
		case cluster2.Name:
			return cluster2Client, nil
		default:
			return nil, fmt.Errorf("Unknown cluster")
		}
	}
	ToFederatedInformerForTestOnly(deploymentController.fedDeploymentInformer).SetClientFactory(clientFactory)
	ToFederatedInformerForTestOnly(deploymentController.fedPodInformer).SetClientFactory(clientFactory)

	stop := make(chan struct{})
	go deploymentController.Run(5, stop)

	// Create deployment. Expect to see it in cluster1.
	dep1 := newDeploymentWithReplicas("depA", 6)
	deploymentsWatch.Add(dep1)
	checkDeployment := func(base *extensionsv1.Deployment, replicas int32) CheckingFunction {
		return func(obj runtime.Object) error {
			if obj == nil {
				return fmt.Errorf("Observed object is nil")
			}
			d := obj.(*extensionsv1.Deployment)
			if err := CompareObjectMeta(base.ObjectMeta, d.ObjectMeta); err != nil {
				return err
			}
			if replicas != *d.Spec.Replicas {
				return fmt.Errorf("Replica count is different expected:%d observed:%d", replicas, *d.Spec.Replicas)
			}
			return nil
		}
	}
	assert.NoError(t, CheckObjectFromChan(cluster1CreateChan, checkDeployment(dep1, *dep1.Spec.Replicas)))
	err := WaitForStoreUpdate(
		deploymentController.fedDeploymentInformer.GetTargetStore(),
		cluster1.Name, types.NamespacedName{Namespace: dep1.Namespace, Name: dep1.Name}.String(), wait.ForeverTestTimeout)
	assert.Nil(t, err, "deployment should have appeared in the informer store")

	// Increase replica count. Expect to see the update in cluster1.
	newRep := int32(8)
	dep1.Spec.Replicas = &newRep
	deploymentsWatch.Modify(dep1)
	assert.NoError(t, CheckObjectFromChan(cluster1UpdateChan, checkDeployment(dep1, *dep1.Spec.Replicas)))

	// Add new cluster. Although rebalance = false, no pods have been created yet so it should
	// rebalance anyway.
	clusterWatch.Add(cluster2)
	assert.NoError(t, CheckObjectFromChan(cluster1UpdateChan, checkDeployment(dep1, *dep1.Spec.Replicas/2)))
	assert.NoError(t, CheckObjectFromChan(cluster2CreateChan, checkDeployment(dep1, *dep1.Spec.Replicas/2)))

	// Add new deployment with non-default replica placement preferences.
	dep2 := newDeploymentWithReplicas("deployment2", 9)
	dep2.Annotations = make(map[string]string)
	dep2.Annotations[FedDeploymentPreferencesAnnotation] = `{"rebalance": true,
		  "clusters": {
		    "cluster1": {"weight": 2},
		    "cluster2": {"weight": 1}
		}}`
	deploymentsWatch.Add(dep2)
	assert.NoError(t, CheckObjectFromChan(cluster1CreateChan, checkDeployment(dep2, 6)))
	assert.NoError(t, CheckObjectFromChan(cluster2CreateChan, checkDeployment(dep2, 3)))
}

func GetDeploymentFromChan(c chan runtime.Object) *extensionsv1.Deployment {
	secret := GetObjectFromChan(c).(*extensionsv1.Deployment)
	return secret
}

func newDeploymentWithReplicas(name string, replicas int32) *extensionsv1.Deployment {
	return &extensionsv1.Deployment{
		ObjectMeta: apiv1.ObjectMeta{
			Name:      name,
			Namespace: apiv1.NamespaceDefault,
			SelfLink:  "/api/v1/namespaces/default/deployments/name",
		},
		Spec: extensionsv1.DeploymentSpec{
			Replicas: &replicas,
		},
	}
}
<|endoftext|>