<commit_before>package ParquetFile

//go:generate mockgen -destination=../mocks/mock_s3.go -package=mocks github.com/aws/aws-sdk-go/service/s3/s3iface S3API

import (
	"context"
	"errors"
	"fmt"
	"io"
	"sync"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3iface"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
)

// S3File is ParquetFile for AWS S3
type S3File struct {
	ctx    context.Context
	client s3iface.S3API
	offset int64
	whence int

	// write-related fields
	writeOpened bool
	writeDone   chan error
	pipeReader  *io.PipeReader
	pipeWriter  *io.PipeWriter
	uploader    *s3manager.Uploader

	// read-related fields
	readOpened bool
	fileSize   int64
	downloader *s3manager.Downloader

	lock       sync.RWMutex
	err        error
	BucketName string
	Key        string
}

const (
	rangeHeader       = "bytes=%d-%d"
	rangeHeaderSuffix = "bytes=%d"
)

var (
	errWhence        = errors.New("Seek: invalid whence")
	errInvalidOffset = errors.New("Seek: invalid offset")
	errFailedUpload  = errors.New("Write: failed upload")
	activeS3Session  *session.Session
	sessLock         sync.RWMutex
)

// NewS3FileWriter creates an S3 FileWriter, to be used with NewParquetWriter
func NewS3FileWriter(ctx context.Context, bucket string, key string, cfgs ...*aws.Config) (ParquetFile, error) {
	sessLock.RLock()
	sess := activeS3Session
	sessLock.RUnlock()
	if sess == nil {
		sessLock.Lock()
		activeS3Session = session.Must(session.NewSession())
		sessLock.Unlock()
	}

	file := &S3File{
		ctx:        ctx,
		client:     s3.New(activeS3Session, cfgs...),
		writeDone:  make(chan error),
		BucketName: bucket,
		Key:        key,
	}

	return file.Create(key)
}

// NewS3FileReader creates an S3 FileReader, to be used with NewParquetReader
func NewS3FileReader(ctx context.Context, bucket string, key string, cfgs ...*aws.Config) (ParquetFile, error) {
	sessLock.RLock()
	sess := activeS3Session
	sessLock.RUnlock()
	if sess == nil {
		sessLock.Lock()
		activeS3Session = session.Must(session.NewSession())
		sessLock.Unlock()
	}

	s3Client := s3.New(activeS3Session, cfgs...)
	s3Downloader := s3manager.NewDownloaderWithClient(s3Client)

	file := &S3File{
		ctx:        ctx,
		client:     s3Client,
		downloader: s3Downloader,
		BucketName: bucket,
		Key:        key,
	}

	return file.Open(key)
}

// Seek tracks the offset for the next Read. Has no effect on Write.
func (s *S3File) Seek(offset int64, whence int) (int64, error) {
	if whence < io.SeekStart || whence > io.SeekEnd {
		return 0, errWhence
	}

	if s.fileSize > 0 {
		switch whence {
		case io.SeekStart:
			if offset < 0 || offset > s.fileSize {
				return 0, errInvalidOffset
			}
		case io.SeekCurrent:
			offset += s.offset
			if offset < 0 || offset > s.fileSize {
				return 0, errInvalidOffset
			}
		case io.SeekEnd:
			if offset > -1 || -offset > s.fileSize {
				return 0, errInvalidOffset
			}
		}
	}

	s.offset = offset
	s.whence = whence
	return s.offset, nil
}

// Read up to len(p) bytes into p and return the number of bytes read
func (s *S3File) Read(p []byte) (n int, err error) {
	if s.fileSize > 0 && s.offset >= s.fileSize {
		return 0, io.EOF
	}

	numBytes := len(p)
	getObjRange := s.getBytesRange(numBytes)
	getObj := &s3.GetObjectInput{
		Bucket: aws.String(s.BucketName),
		Key:    aws.String(s.Key),
	}
	if len(getObjRange) > 0 {
		getObj.Range = aws.String(getObjRange)
	}

	wab := aws.NewWriteAtBuffer(p)
	bytesDownloaded, err := s.downloader.Download(wab, getObj)
	if err != nil {
		return 0, err
	}

	s.offset += bytesDownloaded
	if buf := wab.Bytes(); len(buf) > numBytes {
		// backing buffer reassigned, copy over some of the data
		copy(p, buf)
		bytesDownloaded = int64(len(p))
	}

	return int(bytesDownloaded), err
}

// Write len(p) bytes from p to the S3 data stream
func (s *S3File) Write(p []byte) (n int, err error) {
	s.lock.RLock()
	writeOpened := s.writeOpened
	s.lock.RUnlock()
	if !writeOpened {
		s.openWrite()
	}

	s.lock.RLock()
	writeError := s.err
	s.lock.RUnlock()
	if writeError != nil {
		return 0, writeError
	}

	// prevent further writes upon error
	bytesWritten, writeError := s.pipeWriter.Write(p)
	if writeError != nil {
		s.lock.Lock()
		s.err = writeError
		s.lock.Unlock()

		s.pipeWriter.CloseWithError(err)
		return 0, writeError
	}

	return bytesWritten, nil
}

// Close signals write completion and cleans up any
// open streams. Will block until pending uploads are complete.
func (s *S3File) Close() error {
	var err error

	if s.pipeWriter != nil {
		if err = s.pipeWriter.Close(); err != nil {
			return err
		}
	}

	// wait for pending uploads
	if s.writeDone != nil {
		err = <-s.writeDone
	}

	return err
}

// Open creates a new S3 File instance to perform concurrent reads
func (s *S3File) Open(name string) (ParquetFile, error) {
	s.lock.RLock()
	readOpened := s.readOpened
	s.lock.RUnlock()
	if !readOpened {
		if err := s.openRead(); err != nil {
			return nil, err
		}
	}

	// create a new instance
	pf := &S3File{
		ctx:        s.ctx,
		client:     s.client,
		downloader: s.downloader,
		BucketName: s.BucketName,
		Key:        name,
		readOpened: s.readOpened,
		fileSize:   s.fileSize,
		offset:     0,
	}
	return pf, nil
}

// Create creates a new S3 File instance to perform writes
func (s *S3File) Create(key string) (ParquetFile, error) {
	pf := &S3File{
		ctx:        s.ctx,
		client:     s.client,
		BucketName: s.BucketName,
		Key:        key,
		writeDone:  make(chan error),
	}
	pf.openWrite()
	return pf, nil
}

// openWrite creates an S3 uploader that consumes the Reader end of an io.Pipe.
// Calling Close signals write completion.
func (s *S3File) openWrite() {
	pr, pw := io.Pipe()
	uploader := s3manager.NewUploaderWithClient(s.client)
	s.lock.Lock()
	s.pipeReader = pr
	s.pipeWriter = pw
	s.writeOpened = true
	s.uploader = uploader
	s.lock.Unlock()

	uploadParams := &s3manager.UploadInput{
		Bucket: aws.String(s.BucketName),
		Key:    aws.String(s.Key),
		Body:   s.pipeReader,
	}

	go func(uploader *s3manager.Uploader, params *s3manager.UploadInput, done chan error) {
		defer close(done)

		// upload data and signal done when complete
		_, err := uploader.Upload(params)
		if err != nil {
			s.lock.Lock()
			s.err = err
			s.lock.Unlock()

			if s.writeOpened {
				s.pipeWriter.CloseWithError(err)
			}
		}

		done <- err
	}(s.uploader, uploadParams, s.writeDone)
}

// openRead verifies the requested file is accessible and
// tracks the file size
func (s *S3File) openRead() error {
	hoi := &s3.HeadObjectInput{
		Bucket: aws.String(s.BucketName),
		Key:    aws.String(s.Key),
	}

	hoo, err := s.client.HeadObject(hoi)
	if err != nil {
		return err
	}

	s.lock.Lock()
	s.readOpened = true
	if hoo.ContentLength != nil {
		s.fileSize = *hoo.ContentLength
	}
	s.lock.Unlock()

	return nil
}

// getBytesRange returns the range request header string
func (s *S3File) getBytesRange(numBytes int) string {
	var (
		byteRange string
		begin     int64
		end       int64
	)

	// Processing for unknown file size relies on the requestor to
	// know which ranges are valid. May occur if caller is missing HEAD permissions.
	if s.fileSize < 1 {
		switch s.whence {
		case io.SeekStart, io.SeekCurrent:
			byteRange = fmt.Sprintf(rangeHeader, s.offset, s.offset+int64(numBytes)-1)
		case io.SeekEnd:
			byteRange = fmt.Sprintf(rangeHeaderSuffix, s.offset)
		}
		return byteRange
	}

	switch s.whence {
	case io.SeekStart, io.SeekCurrent:
		begin = s.offset
	case io.SeekEnd:
		begin = s.fileSize + s.offset
	default:
		return byteRange
	}

	endIndex := s.fileSize - 1
	if begin < 0 {
		begin = 0
	}
	end = begin + int64(numBytes) - 1
	if end > endIndex {
		end = endIndex
	}

	byteRange = fmt.Sprintf(rangeHeader, begin, end)
	return byteRange
}
<commit_msg>fix Open name bug, address CR comment<commit_after>package ParquetFile

//go:generate mockgen -destination=../mocks/mock_s3.go -package=mocks github.com/aws/aws-sdk-go/service/s3/s3iface S3API

import (
	"context"
	"errors"
	"fmt"
	"io"
	"sync"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3iface"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
)

// S3File is ParquetFile for AWS S3
type S3File struct {
	ctx    context.Context
	client s3iface.S3API
	offset int64
	whence int

	// write-related fields
	writeOpened bool
	writeDone   chan error
	pipeReader  *io.PipeReader
	pipeWriter  *io.PipeWriter
	uploader    *s3manager.Uploader

	// read-related fields
	readOpened bool
	fileSize   int64
	downloader *s3manager.Downloader

	lock       sync.RWMutex
	err        error
	BucketName string
	Key        string
}

const (
	rangeHeader       = "bytes=%d-%d"
	rangeHeaderSuffix = "bytes=%d"
)

var (
	errWhence        = errors.New("Seek: invalid whence")
	errInvalidOffset = errors.New("Seek: invalid offset")
	errFailedUpload  = errors.New("Write: failed upload")
	activeS3Session  *session.Session
	sessLock         sync.RWMutex
)

// NewS3FileWriter creates an S3 FileWriter, to be used with NewParquetWriter
func NewS3FileWriter(ctx context.Context, bucket string, key string, cfgs ...*aws.Config) (ParquetFile, error) {
	sessLock.RLock()
	sess := activeS3Session
	sessLock.RUnlock()
	if sess == nil {
		sessLock.Lock()
		if activeS3Session == nil {
			activeS3Session = session.Must(session.NewSession())
		}
		sessLock.Unlock()
	}

	file := &S3File{
		ctx:        ctx,
		client:     s3.New(activeS3Session, cfgs...),
		writeDone:  make(chan error),
		BucketName: bucket,
		Key:        key,
	}

	return file.Create(key)
}

// NewS3FileReader creates an S3 FileReader, to be used with NewParquetReader
func NewS3FileReader(ctx context.Context, bucket string, key string, cfgs ...*aws.Config) (ParquetFile, error) {
	sessLock.RLock()
	sess := activeS3Session
	sessLock.RUnlock()
	if sess == nil {
		sessLock.Lock()
		if activeS3Session == nil {
			activeS3Session = session.Must(session.NewSession())
		}
		sessLock.Unlock()
	}

	s3Client := s3.New(activeS3Session, cfgs...)
	s3Downloader := s3manager.NewDownloaderWithClient(s3Client)

	file := &S3File{
		ctx:        ctx,
		client:     s3Client,
		downloader: s3Downloader,
		BucketName: bucket,
		Key:        key,
	}

	return file.Open(key)
}

// Seek tracks the offset for the next Read. Has no effect on Write.
func (s *S3File) Seek(offset int64, whence int) (int64, error) {
	if whence < io.SeekStart || whence > io.SeekEnd {
		return 0, errWhence
	}

	if s.fileSize > 0 {
		switch whence {
		case io.SeekStart:
			if offset < 0 || offset > s.fileSize {
				return 0, errInvalidOffset
			}
		case io.SeekCurrent:
			offset += s.offset
			if offset < 0 || offset > s.fileSize {
				return 0, errInvalidOffset
			}
		case io.SeekEnd:
			if offset > -1 || -offset > s.fileSize {
				return 0, errInvalidOffset
			}
		}
	}

	s.offset = offset
	s.whence = whence
	return s.offset, nil
}

// Read up to len(p) bytes into p and return the number of bytes read
func (s *S3File) Read(p []byte) (n int, err error) {
	if s.fileSize > 0 && s.offset >= s.fileSize {
		return 0, io.EOF
	}

	numBytes := len(p)
	getObjRange := s.getBytesRange(numBytes)
	getObj := &s3.GetObjectInput{
		Bucket: aws.String(s.BucketName),
		Key:    aws.String(s.Key),
	}
	if len(getObjRange) > 0 {
		getObj.Range = aws.String(getObjRange)
	}

	wab := aws.NewWriteAtBuffer(p)
	bytesDownloaded, err := s.downloader.Download(wab, getObj)
	if err != nil {
		return 0, err
	}

	s.offset += bytesDownloaded
	if buf := wab.Bytes(); len(buf) > numBytes {
		// backing buffer reassigned, copy over some of the data
		copy(p, buf)
		bytesDownloaded = int64(len(p))
	}

	return int(bytesDownloaded), err
}

// Write len(p) bytes from p to the S3 data stream
func (s *S3File) Write(p []byte) (n int, err error) {
	s.lock.RLock()
	writeOpened := s.writeOpened
	s.lock.RUnlock()
	if !writeOpened {
		s.openWrite()
	}

	s.lock.RLock()
	writeError := s.err
	s.lock.RUnlock()
	if writeError != nil {
		return 0, writeError
	}

	// prevent further writes upon error
	bytesWritten, writeError := s.pipeWriter.Write(p)
	if writeError != nil {
		s.lock.Lock()
		s.err = writeError
		s.lock.Unlock()

		s.pipeWriter.CloseWithError(err)
		return 0, writeError
	}

	return bytesWritten, nil
}

// Close signals write completion and cleans up any
// open streams. Will block until pending uploads are complete.
func (s *S3File) Close() error {
	var err error

	if s.pipeWriter != nil {
		if err = s.pipeWriter.Close(); err != nil {
			return err
		}
	}

	// wait for pending uploads
	if s.writeDone != nil {
		err = <-s.writeDone
	}

	return err
}

// Open creates a new S3 File instance to perform concurrent reads
func (s *S3File) Open(name string) (ParquetFile, error) {
	s.lock.RLock()
	readOpened := s.readOpened
	s.lock.RUnlock()
	if !readOpened {
		if err := s.openRead(); err != nil {
			return nil, err
		}
	}

	// ColumBuffer passes in an empty string for name
	if len(name) == 0 {
		name = s.Key
	}

	// create a new instance
	pf := &S3File{
		ctx:        s.ctx,
		client:     s.client,
		downloader: s.downloader,
		BucketName: s.BucketName,
		Key:        name,
		readOpened: s.readOpened,
		fileSize:   s.fileSize,
		offset:     0,
	}
	return pf, nil
}

// Create creates a new S3 File instance to perform writes
func (s *S3File) Create(key string) (ParquetFile, error) {
	pf := &S3File{
		ctx:        s.ctx,
		client:     s.client,
		BucketName: s.BucketName,
		Key:        key,
		writeDone:  make(chan error),
	}
	pf.openWrite()
	return pf, nil
}

// openWrite creates an S3 uploader that consumes the Reader end of an io.Pipe.
// Calling Close signals write completion.
func (s *S3File) openWrite() {
	pr, pw := io.Pipe()
	uploader := s3manager.NewUploaderWithClient(s.client)
	s.lock.Lock()
	s.pipeReader = pr
	s.pipeWriter = pw
	s.writeOpened = true
	s.uploader = uploader
	s.lock.Unlock()

	uploadParams := &s3manager.UploadInput{
		Bucket: aws.String(s.BucketName),
		Key:    aws.String(s.Key),
		Body:   s.pipeReader,
	}

	go func(uploader *s3manager.Uploader, params *s3manager.UploadInput, done chan error) {
		defer close(done)

		// upload data and signal done when complete
		_, err := uploader.Upload(params)
		if err != nil {
			s.lock.Lock()
			s.err = err
			s.lock.Unlock()

			if s.writeOpened {
				s.pipeWriter.CloseWithError(err)
			}
		}

		done <- err
	}(s.uploader, uploadParams, s.writeDone)
}

// openRead verifies the requested file is accessible and
// tracks the file size
func (s *S3File) openRead() error {
	hoi := &s3.HeadObjectInput{
		Bucket: aws.String(s.BucketName),
		Key:    aws.String(s.Key),
	}

	hoo, err := s.client.HeadObject(hoi)
	if err != nil {
		return err
	}

	s.lock.Lock()
	s.readOpened = true
	if hoo.ContentLength != nil {
		s.fileSize = *hoo.ContentLength
	}
	s.lock.Unlock()

	return nil
}

// getBytesRange returns the range request header string
func (s *S3File) getBytesRange(numBytes int) string {
	var (
		byteRange string
		begin     int64
		end       int64
	)

	// Processing for unknown file size relies on the requestor to
	// know which ranges are valid. May occur if caller is missing HEAD permissions.
	if s.fileSize < 1 {
		switch s.whence {
		case io.SeekStart, io.SeekCurrent:
			byteRange = fmt.Sprintf(rangeHeader, s.offset, s.offset+int64(numBytes)-1)
		case io.SeekEnd:
			byteRange = fmt.Sprintf(rangeHeaderSuffix, s.offset)
		}
		return byteRange
	}

	switch s.whence {
	case io.SeekStart, io.SeekCurrent:
		begin = s.offset
	case io.SeekEnd:
		begin = s.fileSize + s.offset
	default:
		return byteRange
	}

	endIndex := s.fileSize - 1
	if begin < 0 {
		begin = 0
	}
	end = begin + int64(numBytes) - 1
	if end > endIndex {
		end = endIndex
	}

	byteRange = fmt.Sprintf(rangeHeader, begin, end)
	return byteRange
}
<|endoftext|>


<commit_before>package elasticsearch

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"

	"github.com/elastic/libbeat/logp"
)

type Elasticsearch struct {
	MaxRetries     int
	connectionPool ConnectionPool
	client         *http.Client
}

type QueryResult struct {
	Ok      bool            `json:"ok"`
	Index   string          `json:"_index"`
	Type    string          `json:"_type"`
	Id      string          `json:"_id"`
	Source  json.RawMessage `json:"_source"`
	Version int             `json:"_version"`
	Found   bool            `json:"found"`
	Exists  bool            `json:"exists"`
	Created bool            `json:"created"`
	Matches []string        `json:"matches"`
}

type SearchResults struct {
	Took   int                        `json:"took"`
	Shards json.RawMessage            `json:"_shards"`
	Hits   Hits                       `json:"hits"`
	Aggs   map[string]json.RawMessage `json:"aggregations"`
}

type Hits struct {
	Total int
	Hits  []json.RawMessage `json:"hits"`
}

func (r QueryResult) String() string {
	out, err := json.Marshal(r)
	if err != nil {
		return "ERROR"
	}
	return string(out)
}

const (
	default_max_retries = 3
)

// Create a connection to Elasticsearch
func NewElasticsearch(urls []string, username string, password string) *Elasticsearch {

	var connection_pool ConnectionPool
	connection_pool.SetConnections(urls, username, password)

	es := Elasticsearch{
		connectionPool: connection_pool,
		client:         &http.Client{},
		MaxRetries:     default_max_retries,
	}
	return &es
}

// Encode parameters in url
func UrlEncode(params map[string]string) string {
	var values url.Values = url.Values{}

	for key, val := range params {
		values.Add(key, string(val))
	}
	return values.Encode()
}

// Create path out of index, doc_type and id that is used for querying Elasticsearch
func MakePath(index string, doc_type string, id string) (string, error) {

	var path string
	if len(doc_type) > 0 {
		if len(id) > 0 {
			path = fmt.Sprintf("/%s/%s/%s", index, doc_type, id)
		} else {
			path = fmt.Sprintf("/%s/%s", index, doc_type)
		}
	} else {
		if len(id) > 0 {
			if len(index) > 0 {
				path = fmt.Sprintf("/%s/%s", index, id)
			} else {
				path = fmt.Sprintf("/%s", id)
			}
		} else {
			path = fmt.Sprintf("/%s", index)
		}
	}
	return path, nil
}

func ReadQueryResult(obj []byte) (*QueryResult, error) {

	var result QueryResult
	if obj == nil {
		return nil, nil
	}
	err := json.Unmarshal(obj, &result)
	if err != nil {
		return nil, err
	}
	return &result, err
}

func ReadSearchResult(obj []byte) (*SearchResults, error) {

	var result SearchResults
	if obj == nil {
		return nil, nil
	}
	err := json.Unmarshal(obj, &result)
	if err != nil {
		return nil, err
	}
	return &result, err
}

func (es *Elasticsearch) SetMaxRetries(max_retries int) {
	es.MaxRetries = max_retries
}

// Perform the actual request. If the operation was successful, mark it as live and return the response.
// If it fails, mark it as dead for a period of time.
// It returns the response, if it should retry sending the request and the error
func (es *Elasticsearch) PerformRequest(conn *Connection, req *http.Request) ([]byte, bool, error) {

	req.Header.Add("Accept", "application/json")
	if conn.Username != "" || conn.Password != "" {
		req.SetBasicAuth(conn.Username, conn.Password)
	}

	resp, err := es.client.Do(req)
	if err != nil {
		// request fails
		es.connectionPool.MarkDead(conn)
		return nil, true, fmt.Errorf("Sending the request fails: %s", err)
	}

	if resp.StatusCode > 299 {
		// request fails
		if resp.StatusCode == http.StatusServiceUnavailable ||
			resp.StatusCode == http.StatusGatewayTimeout {
			// status code in {503, 504}
			es.connectionPool.MarkDead(conn)
		}
		return nil, false, fmt.Errorf("%v", resp.Status)
	}

	defer resp.Body.Close()
	obj, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		es.connectionPool.MarkDead(conn)
		return nil, true, fmt.Errorf("Reading the response fails: %s", err)
	}

	// request with success
	es.connectionPool.MarkLive(conn)

	return obj, false, nil

}

// Create an HTTP request and send it to Elasticsearch. The request is retransmitted max_retries
// before returning an error.
func (es *Elasticsearch) Request(method string, path string,
	params map[string]string, body interface{}) ([]byte, error) {

	var errors []error

	for attempt := 0; attempt < es.MaxRetries; attempt++ {

		conn := es.connectionPool.GetConnection()
		logp.Debug("elasticsearch", "Use connection %s", conn.Url)

		url := conn.Url + path
		if len(params) > 0 {
			url = url + "?" + UrlEncode(params)
		}

		logp.Debug("elasticsearch", "%s %s %s", method, url, body)

		var obj []byte
		var err error
		if body != nil {
			obj, err = json.Marshal(body)
			if err != nil {
				return nil, fmt.Errorf("Fail to JSON encode the body: %s", err)
			}
		} else {
			obj = nil
		}
		req, err := http.NewRequest(method, url, bytes.NewReader(obj))
		if err != nil {
			return nil, fmt.Errorf("NewRequest fails: %s", err)
		}

		resp, retry, err := es.PerformRequest(conn, req)
		if retry == true {
			// retry
			if err != nil {
				errors = append(errors, err)
			}
			continue
		}
		if err != nil {
			return nil, err
		}
		return resp, nil

	}

	logp.Warn("Request fails to be send after %d retries", es.MaxRetries)

	return nil, fmt.Errorf("Request fails after %d retries. Errors: %v", es.MaxRetries, errors)
}

// Index adds or updates a typed JSON document in a specified index, making it
// searchable. In case id is empty, a new id is created over a HTTP POST request.
// Otherwise, a HTTP PUT request is issued.
// Implements: http://www.elastic.co/guide/en/elasticsearch/reference/current/docs-index_.html
func (es *Elasticsearch) Index(index string, doc_type string, id string,
	params map[string]string, body interface{}) (*QueryResult, error) {

	var method string

	path, err := MakePath(index, doc_type, id)
	if err != nil {
		return nil, fmt.Errorf("MakePath fails: %s", err)
	}
	if len(id) == 0 {
		method = "POST"
	} else {
		method = "PUT"
	}
	resp, err := es.Request(method, path, params, body)
	if err != nil {
		return nil, err
	}
	return ReadQueryResult(resp)
}

// Refresh an index. Call this after doing inserts or creating/deleting
// indexes in unit tests.
func (es *Elasticsearch) Refresh(index string) (*QueryResult, error) {
	path, err := MakePath(index, "", "_refresh")
	if err != nil {
		return nil, err
	}
	resp, err := es.Request("POST", path, nil, nil)
	if err != nil {
		return nil, err
	}

	return ReadQueryResult(resp)
}

// Adds a typed JSON document in a specific index, making it searchable.
// Implements: <http://www.elastic.co/guide/en/elasticsearch/reference/current/docs-index_.html>
func (es *Elasticsearch) CreateIndex(index string) (*QueryResult, error) {

	path, err := MakePath(index, "", "")
	if err != nil {
		return nil, err
	}

	resp, err := es.Request("PUT", path, nil, nil)
	if err != nil {
		return nil, err
	}

	return ReadQueryResult(resp)
}

// Deletes a typed JSON document from a specific index based on its id.
// Implements: http://www.elastic.co/guide/en/elasticsearch/reference/current/docs-delete.html
func (es *Elasticsearch) Delete(index string, doc_type string, id string, params map[string]string) (*QueryResult, error) {

	path, err := MakePath(index, doc_type, id)
	if err != nil {
		return nil, err
	}

	resp, err := es.Request("DELETE", path, params, nil)
	if err != nil {
		return nil, err
	}

	return ReadQueryResult(resp)
}

// A search request can be executed purely using a URI by providing request parameters.
// Implements: http://www.elastic.co/guide/en/elasticsearch/reference/current/search-uri-request.html
func (es *Elasticsearch) SearchUri(index string, doc_type string, params map[string]string) (*SearchResults, error) {

	path, err := MakePath(index, doc_type, "_search")
	if err != nil {
		return nil, err
	}

	resp, err := es.Request("GET", path, params, nil)
	if err != nil {
		return nil, err
	}
	return ReadSearchResult(resp)
}
<commit_msg>Change when to retry sending the http request to a different node<commit_after>package elasticsearch

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"

	"github.com/elastic/libbeat/logp"
)

type Elasticsearch struct {
	MaxRetries     int
	connectionPool ConnectionPool
	client         *http.Client
}

type QueryResult struct {
	Ok      bool            `json:"ok"`
	Index   string          `json:"_index"`
	Type    string          `json:"_type"`
	Id      string          `json:"_id"`
	Source  json.RawMessage `json:"_source"`
	Version int             `json:"_version"`
	Found   bool            `json:"found"`
	Exists  bool            `json:"exists"`
	Created bool            `json:"created"`
	Matches []string        `json:"matches"`
}

type SearchResults struct {
	Took   int                        `json:"took"`
	Shards json.RawMessage            `json:"_shards"`
	Hits   Hits                       `json:"hits"`
	Aggs   map[string]json.RawMessage `json:"aggregations"`
}

type Hits struct {
	Total int
	Hits  []json.RawMessage `json:"hits"`
}

func (r QueryResult) String() string {
	out, err := json.Marshal(r)
	if err != nil {
		return "ERROR"
	}
	return string(out)
}

const (
	default_max_retries = 3
)

// Create a connection to Elasticsearch
func NewElasticsearch(urls []string, username string, password string) *Elasticsearch {

	var connection_pool ConnectionPool
	connection_pool.SetConnections(urls, username, password)

	es := Elasticsearch{
		connectionPool: connection_pool,
		client:         &http.Client{},
		MaxRetries:     default_max_retries,
	}
	return &es
}

// Encode parameters in url
func UrlEncode(params map[string]string) string {
	var values url.Values = url.Values{}

	for key, val := range params {
		values.Add(key, string(val))
	}
	return values.Encode()
}

// Create path out of index, doc_type and id that is used for querying Elasticsearch
func MakePath(index string, doc_type string, id string) (string, error) {

	var path string
	if len(doc_type) > 0 {
		if len(id) > 0 {
			path = fmt.Sprintf("/%s/%s/%s", index, doc_type, id)
		} else {
			path = fmt.Sprintf("/%s/%s", index, doc_type)
		}
	} else {
		if len(id) > 0 {
			if len(index) > 0 {
				path = fmt.Sprintf("/%s/%s", index, id)
			} else {
				path = fmt.Sprintf("/%s", id)
			}
		} else {
			path = fmt.Sprintf("/%s", index)
		}
	}
	return path, nil
}

func ReadQueryResult(obj []byte) (*QueryResult, error) {

	var result QueryResult
	if obj == nil {
		return nil, nil
	}
	err := json.Unmarshal(obj, &result)
	if err != nil {
		return nil, err
	}
	return &result, err
}

func ReadSearchResult(obj []byte) (*SearchResults, error) {

	var result SearchResults
	if obj == nil {
		return nil, nil
	}
	err := json.Unmarshal(obj, &result)
	if err != nil {
		return nil, err
	}
	return &result, err
}

func (es *Elasticsearch) SetMaxRetries(max_retries int) {
	es.MaxRetries = max_retries
}

func isConnTimeout(err error) bool {
	return strings.Contains(err.Error(), "i/o timeout")
}

func isConnRefused(err error) bool {
	return strings.Contains(err.Error(), "connection refused")
}

// Perform the actual request. If the operation was successful, mark it as live and return the response.
// Mark the Elasticsearch node as dead for a period of time in the case the http request fails with Connection
// Timeout, Connection Refused or returns one of the 503,504 Error Replies.
// It returns the response, if it should retry sending the request and the error
func (es *Elasticsearch) PerformRequest(conn *Connection, req *http.Request) ([]byte, bool, error) {

	req.Header.Add("Accept", "application/json")
	if conn.Username != "" || conn.Password != "" {
		req.SetBasicAuth(conn.Username, conn.Password)
	}

	resp, err := es.client.Do(req)
	if err != nil {
		// request fails
		do_retry := false
		if isConnTimeout(err) || isConnRefused(err) {
			es.connectionPool.MarkDead(conn)
			do_retry = true
		}
		return nil, do_retry, fmt.Errorf("Sending the request fails: %s", err)
	}

	if resp.StatusCode > 299 {
		// request fails
		do_retry := false
		if resp.StatusCode == http.StatusServiceUnavailable ||
			resp.StatusCode == http.StatusGatewayTimeout {
			// status code in {503, 504}
			es.connectionPool.MarkDead(conn)
			do_retry = true
		}
		return nil, do_retry, fmt.Errorf("%v", resp.Status)
	}

	defer resp.Body.Close()
	obj, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		es.connectionPool.MarkDead(conn)
		return nil, true, fmt.Errorf("Reading the response fails: %s", err)
	}

	// request with success
	es.connectionPool.MarkLive(conn)

	return obj, false, nil

}

// Create an HTTP request and send it to Elasticsearch. The request is retransmitted max_retries
// before returning an error.
func (es *Elasticsearch) Request(method string, path string,
	params map[string]string, body interface{}) ([]byte, error) {

	var errors []error

	for attempt := 0; attempt < es.MaxRetries; attempt++ {

		conn := es.connectionPool.GetConnection()
		logp.Debug("elasticsearch", "Use connection %s", conn.Url)

		url := conn.Url + path
		if len(params) > 0 {
			url = url + "?" + UrlEncode(params)
		}

		logp.Debug("elasticsearch", "%s %s %s", method, url, body)

		var obj []byte
		var err error
		if body != nil {
			obj, err = json.Marshal(body)
			if err != nil {
				return nil, fmt.Errorf("Fail to JSON encode the body: %s", err)
			}
		} else {
			obj = nil
		}
		req, err := http.NewRequest(method, url, bytes.NewReader(obj))
		if err != nil {
			return nil, fmt.Errorf("NewRequest fails: %s", err)
		}

		resp, retry, err := es.PerformRequest(conn, req)
		if retry == true {
			// retry
			if err != nil {
				errors = append(errors, err)
			}
			continue
		}
		if err != nil {
			return nil, err
		}
		return resp, nil

	}

	logp.Warn("Request fails to be send after %d retries", es.MaxRetries)

	return nil, fmt.Errorf("Request fails after %d retries. Errors: %v", es.MaxRetries, errors)
}

// Index adds or updates a typed JSON document in a specified index, making it
// searchable. In case id is empty, a new id is created over a HTTP POST request.
// Otherwise, a HTTP PUT request is issued.
// Implements: http://www.elastic.co/guide/en/elasticsearch/reference/current/docs-index_.html
func (es *Elasticsearch) Index(index string, doc_type string, id string,
	params map[string]string, body interface{}) (*QueryResult, error) {

	var method string

	path, err := MakePath(index, doc_type, id)
	if err != nil {
		return nil, fmt.Errorf("MakePath fails: %s", err)
	}
	if len(id) == 0 {
		method = "POST"
	} else {
		method = "PUT"
	}
	resp, err := es.Request(method, path, params, body)
	if err != nil {
		return nil, err
	}
	return ReadQueryResult(resp)
}

// Refresh an index. Call this after doing inserts or creating/deleting
// indexes in unit tests.
func (es *Elasticsearch) Refresh(index string) (*QueryResult, error) {
	path, err := MakePath(index, "", "_refresh")
	if err != nil {
		return nil, err
	}
	resp, err := es.Request("POST", path, nil, nil)
	if err != nil {
		return nil, err
	}

	return ReadQueryResult(resp)
}

// Adds a typed JSON document in a specific index, making it searchable.
// Implements: <http://www.elastic.co/guide/en/elasticsearch/reference/current/docs-index_.html>
func (es *Elasticsearch) CreateIndex(index string) (*QueryResult, error) {

	path, err := MakePath(index, "", "")
	if err != nil {
		return nil, err
	}

	resp, err := es.Request("PUT", path, nil, nil)
	if err != nil {
		return nil, err
	}

	return ReadQueryResult(resp)
}

// Deletes a typed JSON document from a specific index based on its id.
// Implements: http://www.elastic.co/guide/en/elasticsearch/reference/current/docs-delete.html
func (es *Elasticsearch) Delete(index string, doc_type string, id string, params map[string]string) (*QueryResult, error) {

	path, err := MakePath(index, doc_type, id)
	if err != nil {
		return nil, err
	}

	resp, err := es.Request("DELETE", path, params, nil)
	if err != nil {
		return nil, err
	}

	return ReadQueryResult(resp)
}

// A search request can be executed purely using a URI by providing request parameters.
// Implements: http://www.elastic.co/guide/en/elasticsearch/reference/current/search-uri-request.html
func (es *Elasticsearch) SearchUri(index string, doc_type string, params map[string]string) (*SearchResults, error) {

	path, err := MakePath(index, doc_type, "_search")
	if err != nil {
		return nil, err
	}

	resp, err := es.Request("GET", path, params, nil)
	if err != nil {
		return nil, err
	}
	return ReadSearchResult(resp)
}
<|endoftext|>


<commit_before>package logstash

// logstash.go defines the logtash plugin (using lumberjack protocol) as being registered with all
// output plugins

import (
	"crypto/tls"
	"errors"
	"fmt"
	"time"

	"github.com/elastic/libbeat/common"
	"github.com/elastic/libbeat/logp"
	"github.com/elastic/libbeat/outputs"
	"github.com/elastic/libbeat/outputs/mode"
)

var debug = logp.MakeDebug("logstash")

func init() {
	outputs.RegisterOutputPlugin("logstash", logstashOutputPlugin{})
}

type logstashOutputPlugin struct{}

func (p logstashOutputPlugin) NewOutput(
	beat string,
	config *outputs.MothershipConfig,
	topologyExpire int,
) (outputs.Outputer, error) {
	output := &logstash{}
	err := output.init(beat, *config, topologyExpire)
	if err != nil {
		return nil, err
	}
	return output, nil
}

type logstash struct {
	mode  mode.ConnectionMode
	index string
}

const (
	logstashDefaultPort = 10200

	logstashDefaultTimeout = 30 * time.Second
	defaultSendRetries     = 3
)

// ErrNoHostsConfigured indicates missing host or hosts configuration
var ErrNoHostsConfigured = errors.New("no host configuration found")

var waitRetry = time.Duration(1) * time.Second

func (lj *logstash) init(
	beat string,
	config outputs.MothershipConfig,
	topologyExpire int,
) error {
	useTLS := false
	if config.TLS != nil {
		useTLS = !config.TLS.Disabled
	}

	timeout := logstashDefaultTimeout
	if config.Timeout != 0 {
		timeout = time.Duration(config.Timeout) * time.Second
	}

	defaultPort := logstashDefaultPort
	if config.Port != 0 {
		defaultPort = config.Port
	}

	var clients []mode.ProtocolClient
	var err error
	if useTLS {
		var tlsConfig *tls.Config
		tlsConfig, err = outputs.LoadTLSConfig(config.TLS)
		if err != nil {
			return err
		}

		clients, err = makeClients(config, timeout,
			func(host string) (TransportClient, error) {
				return newTLSClient(host, defaultPort, tlsConfig)
			})
	} else {
		clients, err = makeClients(config, timeout,
			func(host string) (TransportClient, error) {
				return newTCPClient(host, defaultPort)
			})
	}
	if err != nil {
		return err
	}

	sendRetries := defaultSendRetries
	if config.Max_retries != nil {
		sendRetries = *config.Max_retries
	}

	var m mode.ConnectionMode
	if len(clients) == 1 {
		m, err = mode.NewSingleConnectionMode(clients[0],
			sendRetries, waitRetry, timeout)
	} else {
		loadBalance := config.LoadBalance == nil || *config.LoadBalance
		if loadBalance {
			m, err = mode.NewLoadBalancerMode(clients, sendRetries, waitRetry, timeout)
		} else {
			m, err = mode.NewFailOverConnectionMode(clients, sendRetries, waitRetry, timeout)
		}
	}
	if err != nil {
		return err
	}

	lj.mode = m
	if config.Index != "" {
		lj.index = config.Index
	} else {
		lj.index = beat
	}
	return nil
}

func makeClients(
	config outputs.MothershipConfig,
	timeout time.Duration,
	newTransp func(string) (TransportClient, error),
) ([]mode.ProtocolClient, error) {
	switch {
	case len(config.Hosts) > 0:
		var clients []mode.ProtocolClient
		for _, host := range config.Hosts {
			transp, err := newTransp(host)
			if err != nil {
				for _, client := range clients {
					_ = client.Close() // ignore error
				}
				return nil, err
			}
			client := newLumberjackClient(transp, timeout)
			clients = append(clients, client)
		}
		return clients, nil
	case config.Host != "":
		transp, err := newTransp(config.Host)
		if err != nil {
			return nil, err
		}
		return []mode.ProtocolClient{newLumberjackClient(transp, timeout)}, nil
	default:
		return nil, ErrNoHostsConfigured
	}
}

// TODO: update Outputer interface to support multiple events for batch-like
//       processing (e.g. for filebeat). Batch like processing might reduce
//       send/receive overhead per event for other implementors too.
func (lj *logstash) PublishEvent(
	signaler outputs.Signaler,
	ts time.Time,
	event common.MapStr,
) error {
	lj.addMeta(event)
	return lj.mode.PublishEvent(signaler, event)
}

// BulkPublish implements the BulkOutputer interface pushing a bulk of events
// via lumberjack.
func (lj *logstash) BulkPublish(
	trans outputs.Signaler,
	ts time.Time,
	events []common.MapStr,
) error {
	for _, event := range events {
		lj.addMeta(event)
	}
	return lj.mode.PublishEvents(trans, events)
}

// addMeta adapts events to be compatible with logstash forwarer messages by renaming
// the "message" field to "line". The lumberjack server in logstash will
// decode/rename the "line" field into "message".
func (lj *logstash) addMeta(event common.MapStr) {
	// add metadata for indexing
	ts := time.Time(event["timestamp"].(common.Time)).UTC()
	index := fmt.Sprintf("%s-%02d.%02d.%02d", lj.index,
		ts.Year(), ts.Month(), ts.Day())
	event["@metadata"] = common.MapStr{
		"index": index,
		"type":  event["type"].(string),
	}
}
<commit_msg>logstash load balancing must be explicitely set<commit_after>package logstash

// logstash.go defines the logtash plugin (using lumberjack protocol) as being registered with all
// output plugins

import (
	"crypto/tls"
	"errors"
	"fmt"
	"time"

	"github.com/elastic/libbeat/common"
	"github.com/elastic/libbeat/logp"
	"github.com/elastic/libbeat/outputs"
	"github.com/elastic/libbeat/outputs/mode"
)

var debug = logp.MakeDebug("logstash")

func init() {
	outputs.RegisterOutputPlugin("logstash", logstashOutputPlugin{})
}

type logstashOutputPlugin struct{}

func (p logstashOutputPlugin) NewOutput(
	beat string,
	config *outputs.MothershipConfig,
	topologyExpire int,
) (outputs.Outputer, error) {
	output := &logstash{}
	err := output.init(beat, *config, topologyExpire)
	if err != nil {
		return nil, err
	}
	return output, nil
}

type logstash struct {
	mode  mode.ConnectionMode
	index string
}

const (
	logstashDefaultPort = 10200

	logstashDefaultTimeout = 30 * time.Second
	defaultSendRetries     = 3
)

// ErrNoHostsConfigured indicates missing host or hosts configuration
var ErrNoHostsConfigured = errors.New("no host configuration found")

var waitRetry = time.Duration(1) * time.Second

func (lj *logstash) init(
	beat string,
	config outputs.MothershipConfig,
	topologyExpire int,
) error {
	useTLS := false
	if config.TLS != nil {
		useTLS = !config.TLS.Disabled
	}

	timeout := logstashDefaultTimeout
	if config.Timeout != 0 {
		timeout = time.Duration(config.Timeout) * time.Second
	}

	defaultPort := logstashDefaultPort
	if config.Port != 0 {
		defaultPort = config.Port
	}

	var clients []mode.ProtocolClient
	var err error
	if useTLS {
		var tlsConfig *tls.Config
		tlsConfig, err = outputs.LoadTLSConfig(config.TLS)
		if err != nil {
			return err
		}

		clients, err = makeClients(config, timeout,
			func(host string) (TransportClient, error) {
				return newTLSClient(host, defaultPort, tlsConfig)
			})
	} else {
		clients, err = makeClients(config, timeout,
			func(host string) (TransportClient, error) {
				return newTCPClient(host, defaultPort)
			})
	}
	if err != nil {
		return err
	}

	sendRetries := defaultSendRetries
	if config.Max_retries != nil {
		sendRetries = *config.Max_retries
	}

	var m mode.ConnectionMode
	if len(clients) == 1 {
		m, err = mode.NewSingleConnectionMode(clients[0],
			sendRetries, waitRetry, timeout)
	} else {
		loadBalance := config.LoadBalance != nil && *config.LoadBalance
		if loadBalance {
			m, err = mode.NewLoadBalancerMode(clients, sendRetries, waitRetry, timeout)
		} else {
			m, err = mode.NewFailOverConnectionMode(clients, sendRetries, waitRetry, timeout)
		}
	}
	if err != nil {
		return err
	}

	lj.mode = m
	if config.Index != "" {
		lj.index = config.Index
	} else {
		lj.index = beat
	}
	return nil
}

func makeClients(
	config outputs.MothershipConfig,
	timeout time.Duration,
	newTransp func(string) (TransportClient, error),
) ([]mode.ProtocolClient, error) {
	switch {
	case len(config.Hosts) > 0:
		var clients []mode.ProtocolClient
		for _, host := range config.Hosts {
			transp, err := newTransp(host)
			if err != nil {
				for _, client := range clients {
					_ = client.Close() // ignore error
				}
				return nil, err
			}
			client := newLumberjackClient(transp, timeout)
			clients = append(clients, client)
		}
		return clients, nil
	case config.Host != "":
		transp, err := newTransp(config.Host)
		if err != nil {
			return nil, err
		}
		return []mode.ProtocolClient{newLumberjackClient(transp, timeout)}, nil
	default:
		return nil, ErrNoHostsConfigured
	}
}

// TODO: update Outputer interface to support multiple events for batch-like
//       processing (e.g. for filebeat). Batch like processing might reduce
//       send/receive overhead per event for other implementors too.
func (lj *logstash) PublishEvent(
	signaler outputs.Signaler,
	ts time.Time,
	event common.MapStr,
) error {
	lj.addMeta(event)
	return lj.mode.PublishEvent(signaler, event)
}

// BulkPublish implements the BulkOutputer interface pushing a bulk of events
// via lumberjack.
func (lj *logstash) BulkPublish(
	trans outputs.Signaler,
	ts time.Time,
	events []common.MapStr,
) error {
	for _, event := range events {
		lj.addMeta(event)
	}
	return lj.mode.PublishEvents(trans, events)
}

// addMeta adapts events to be compatible with logstash forwarer messages by renaming
// the "message" field to "line". The lumberjack server in logstash will
// decode/rename the "line" field into "message".
func (lj *logstash) addMeta(event common.MapStr) {
	// add metadata for indexing
	ts := time.Time(event["timestamp"].(common.Time)).UTC()
	index := fmt.Sprintf("%s-%02d.%02d.%02d", lj.index,
		ts.Year(), ts.Month(), ts.Day())
	event["@metadata"] = common.MapStr{
		"index": index,
		"type":  event["type"].(string),
	}
}
<|endoftext|>


<commit_before>package fuse

import (
	"bytes"
	"fmt"
	"log"
	"strings"
	"time"
	"unsafe"
)

var sizeOfOutHeader = unsafe.Sizeof(OutHeader{})
var zeroOutBuf [outputHeaderSize]byte

type request struct {
	inputBuf []byte

	// These split up inputBuf.
	inHeader *InHeader      // generic header
	inData   unsafe.Pointer // per op data
	arg      []byte         // flat data.

	filenames []string // filename arguments

	// Unstructured data, a pointer to the relevant XxxxOut struct.
	outData  unsafe.Pointer
	status   Status
	flatData []byte
	fdData   *readResultFd

	// In case of read, keep read result here so we can call
	// Done() on it.
	readResult ReadResult

	// Start timestamp for timing info.
	startTime time.Time

	// All information pertaining to opcode of this request.
	handler *operationHandler

	// Request storage. For large inputs and outputs, use data
	// obtained through bufferpool.
	bufferPoolInputBuf  []byte
	bufferPoolOutputBuf []byte

	// For small pieces of data, we use the following inlines
	// arrays:
	//
	// Output header and structured data.
	outBuf [outputHeaderSize]byte

	// Input, if small enough to fit here.
	smallInputBuf [128]byte

	context Context
}

func (r *request) clear() {
	r.inputBuf = nil
	r.inHeader = nil
	r.inData = nil
	r.arg = nil
	r.filenames = nil
	r.outData = nil
	r.status = OK
	r.flatData = nil
	r.fdData = nil
	r.startTime = time.Time{}
	r.handler = nil
	r.readResult = nil
}

func (r *request) InputDebug() string {
	val := " "
	if r.handler.DecodeIn != nil {
		val = fmt.Sprintf(" data: %v ", Print(r.handler.DecodeIn(r.inData)))
	}

	names := ""
	if r.filenames != nil {
		names = fmt.Sprintf("names: %v", r.filenames)
	}

	if len(r.arg) > 0 {
		names += fmt.Sprintf(" %d bytes", len(r.arg))
	}

	return fmt.Sprintf("Dispatch: %s, NodeId: %v.%v%v",
		operationName(r.inHeader.Opcode), r.inHeader.NodeId, val, names)
}

func (r *request) OutputDebug() string {
	var dataStr string
	if r.handler.DecodeOut != nil && r.outData != nil {
		dataStr = Print(r.handler.DecodeOut(r.outData))
	}

	max := 1024
	if len(dataStr) > max {
		dataStr = dataStr[:max] + fmt.Sprintf(" ...trimmed")
	}

	flatStr := ""
	if r.flatDataSize() > 0 {
		if r.handler.FileNameOut {
			s := strings.TrimRight(string(r.flatData), "\x00")
			flatStr = fmt.Sprintf(" %q", s)
		} else {
			spl := ""
			if r.fdData != nil {
				spl = " (fd data)"
			}
			flatStr = fmt.Sprintf(" %d bytes data%s\n", r.flatDataSize(), spl)
		}
	}

	return fmt.Sprintf("Serialize: %s code: %v value: %v%v",
		operationName(r.inHeader.Opcode), r.status, dataStr, flatStr)
}

// setInput returns true if it takes ownership of the argument, false if not.
func (r *request) setInput(input []byte) bool {
	if len(input) < len(r.smallInputBuf) {
		copy(r.smallInputBuf[:], input)
		r.inputBuf = r.smallInputBuf[:len(input)]
		return false
	}
	r.inputBuf = input
	r.bufferPoolInputBuf = input[:cap(input)]

	return true
}

func (r *request) parse() {
	inHSize := int(unsafe.Sizeof(InHeader{}))
	if len(r.inputBuf) < inHSize {
		log.Printf("Short read for input header: %v", r.inputBuf)
		return
	}

	r.inHeader = (*InHeader)(unsafe.Pointer(&r.inputBuf[0]))
	r.arg = r.inputBuf[:]

	r.handler = getHandler(r.inHeader.Opcode)
	if r.handler == nil {
		log.Printf("Unknown opcode %d", r.inHeader.Opcode)
		r.status = ENOSYS
		return
	}

	if len(r.arg) < int(r.handler.InputSize) {
		log.Printf("Short read for %v: %v", operationName(r.inHeader.Opcode), r.arg)
		r.status = EIO
		return
	}

	if r.handler.InputSize > 0 {
		r.inData = unsafe.Pointer(&r.arg[0])
		r.arg = r.arg[r.handler.InputSize:]
	} else {
		r.arg = r.arg[inHSize:]
	}

	count := r.handler.FileNames
	if count > 0 {
		if count == 1 {
			r.filenames = []string{string(r.arg[:len(r.arg)-1])}
		} else {
			names := bytes.SplitN(r.arg[:len(r.arg)-1], []byte{0}, count)
			r.filenames = make([]string, len(names))
			for i, n := range names {
				r.filenames[i] = string(n)
			}
			if len(names) != count {
				log.Println("filename argument mismatch", names, count)
				r.status = EIO
			}
		}
	}

	copy(r.outBuf[:r.handler.OutputSize+sizeOfOutHeader],
		zeroOutBuf[:r.handler.OutputSize+sizeOfOutHeader])
	r.outData = unsafe.Pointer(&r.outBuf[sizeOfOutHeader])
}

func (r *request) serializeHeader(dataSize int) (header []byte) {
	dataLength := r.handler.OutputSize
	if r.outData == nil || r.status > OK {
		dataLength = 0
	}

	sizeOfOutHeader := unsafe.Sizeof(OutHeader{})
	header = r.outBuf[:sizeOfOutHeader+dataLength]
	o := (*OutHeader)(unsafe.Pointer(&header[0]))
	o.Unique = r.inHeader.Unique
	o.Status = int32(-r.status)
	o.Length = uint32(
		int(sizeOfOutHeader) + int(dataLength) + dataSize)

	var asSlice []byte
	toSlice(&asSlice, r.outData, dataLength)
	copy(header[sizeOfOutHeader:], asSlice)
	return header
}

func (r *request) flatDataSize() int {
	if r.fdData != nil {
		return r.fdData.Size()
	}
	return len(r.flatData)
}
<commit_msg>fuse: show inHeader.Unique in request debug output<commit_after>package fuse

import (
	"bytes"
	"fmt"
	"log"
	"strings"
	"time"
	"unsafe"
)

var sizeOfOutHeader = unsafe.Sizeof(OutHeader{})
var zeroOutBuf [outputHeaderSize]byte

type request struct {
	inputBuf []byte

	// These split up inputBuf.
	inHeader *InHeader      // generic header
	inData   unsafe.Pointer // per op data
	arg      []byte         // flat data.

	filenames []string // filename arguments

	// Unstructured data, a pointer to the relevant XxxxOut struct.
	outData  unsafe.Pointer
	status   Status
	flatData []byte
	fdData   *readResultFd

	// In case of read, keep read result here so we can call
	// Done() on it.
	readResult ReadResult

	// Start timestamp for timing info.
	startTime time.Time

	// All information pertaining to opcode of this request.
	handler *operationHandler

	// Request storage. For large inputs and outputs, use data
	// obtained through bufferpool.
	bufferPoolInputBuf  []byte
	bufferPoolOutputBuf []byte

	// For small pieces of data, we use the following inlines
	// arrays:
	//
	// Output header and structured data.
	outBuf [outputHeaderSize]byte

	// Input, if small enough to fit here.
	smallInputBuf [128]byte

	context Context
}

func (r *request) clear() {
	r.inputBuf = nil
	r.inHeader = nil
	r.inData = nil
	r.arg = nil
	r.filenames = nil
	r.outData = nil
	r.status = OK
	r.flatData = nil
	r.fdData = nil
	r.startTime = time.Time{}
	r.handler = nil
	r.readResult = nil
}

func (r *request) InputDebug() string {
	val := " "
	if r.handler.DecodeIn != nil {
		val = fmt.Sprintf(" data: %v ", Print(r.handler.DecodeIn(r.inData)))
	}

	names := ""
	if r.filenames != nil {
		names = fmt.Sprintf("names: %v", r.filenames)
	}

	if len(r.arg) > 0 {
		names += fmt.Sprintf(" %d bytes", len(r.arg))
	}

	return fmt.Sprintf("Dispatch %d: %s, NodeId: %v.%v%v",
		r.inHeader.Unique, operationName(r.inHeader.Opcode),
		r.inHeader.NodeId, val, names)
}

func (r *request) OutputDebug() string {
	var dataStr string
	if r.handler.DecodeOut != nil && r.outData != nil {
		dataStr = Print(r.handler.DecodeOut(r.outData))
	}

	max := 1024
	if len(dataStr) > max {
		dataStr = dataStr[:max] + fmt.Sprintf(" ...trimmed")
	}

	flatStr := ""
	if r.flatDataSize() > 0 {
		if r.handler.FileNameOut {
			s := strings.TrimRight(string(r.flatData), "\x00")
			flatStr = fmt.Sprintf(" %q", s)
		} else {
			spl := ""
			if r.fdData != nil {
				spl = " (fd data)"
			}
			flatStr = fmt.Sprintf(" %d bytes data%s\n", r.flatDataSize(), spl)
		}
	}

	return fmt.Sprintf("Serialize %d: %s code: %v value: %v%v",
		r.inHeader.Unique, operationName(r.inHeader.Opcode), r.status,
		dataStr, flatStr)
}

// setInput returns true if it takes ownership of the argument, false if not.
func (r *request) setInput(input []byte) bool {
	if len(input) < len(r.smallInputBuf) {
		copy(r.smallInputBuf[:], input)
		r.inputBuf = r.smallInputBuf[:len(input)]
		return false
	}
	r.inputBuf = input
	r.bufferPoolInputBuf = input[:cap(input)]

	return true
}

func (r *request) parse() {
	inHSize := int(unsafe.Sizeof(InHeader{}))
	if len(r.inputBuf) < inHSize {
		log.Printf("Short read for input header: %v", r.inputBuf)
		return
	}

	r.inHeader = (*InHeader)(unsafe.Pointer(&r.inputBuf[0]))
	r.arg = r.inputBuf[:]

	r.handler = getHandler(r.inHeader.Opcode)
	if r.handler == nil {
		log.Printf("Unknown opcode %d", r.inHeader.Opcode)
		r.status = ENOSYS
		return
	}

	if len(r.arg) < int(r.handler.InputSize) {
		log.Printf("Short read for %v: %v", operationName(r.inHeader.Opcode), r.arg)
		r.status = EIO
		return
	}

	if r.handler.InputSize > 0 {
		r.inData = unsafe.Pointer(&r.arg[0])
		r.arg = r.arg[r.handler.InputSize:]
	} else {
		r.arg = r.arg[inHSize:]
	}

	count := r.handler.FileNames
	if count > 0 {
		if count == 1 {
			r.filenames = []string{string(r.arg[:len(r.arg)-1])}
		} else {
			names := bytes.SplitN(r.arg[:len(r.arg)-1], []byte{0}, count)
			r.filenames = make([]string, len(names))
			for i, n := range names {
				r.filenames[i] = string(n)
			}
			if len(names) != count {
				log.Println("filename argument mismatch", names, count)
				r.status = EIO
			}
		}
	}

	copy(r.outBuf[:r.handler.OutputSize+sizeOfOutHeader],
		zeroOutBuf[:r.handler.OutputSize+sizeOfOutHeader])
	r.outData = unsafe.Pointer(&r.outBuf[sizeOfOutHeader])
}

func (r *request) serializeHeader(dataSize int) (header []byte) {
	dataLength := r.handler.OutputSize
	if r.outData == nil || r.status > OK {
		dataLength = 0
	}

	sizeOfOutHeader := unsafe.Sizeof(OutHeader{})
	header = r.outBuf[:sizeOfOutHeader+dataLength]
	o := (*OutHeader)(unsafe.Pointer(&header[0]))
	o.Unique = r.inHeader.Unique
	o.Status = int32(-r.status)
	o.Length = uint32(
		int(sizeOfOutHeader) + int(dataLength) + dataSize)

	var asSlice []byte
	toSlice(&asSlice, r.outData, dataLength)
	copy(header[sizeOfOutHeader:], asSlice)
	return header
}

func (r *request) flatDataSize() int {
	if r.fdData != nil {
		return r.fdData.Size()
	}
	return len(r.flatData)
}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2017-2018 Alibaba Group Holding Limited
 */
package cli

import (
	"fmt"
	"github.com/aliyun/aliyun-cli/i18n"
)

//
// default help flag
var HelpFlag = &Flag{
	Name:         "help",
	Short:        i18n.T("print help", "打印帮助信息"),
	AssignedMode: AssignedNone,
}

//
// CLI Command Context
type Context struct {
	help         bool
	flags        *FlagSet
	unknownFlags *FlagSet
	command      *Command
	completion   *Completion
}

func NewCommandContext() *Context {
	return &Context{
		flags:        NewFlagSet(),
		unknownFlags: nil,
	}
}

func (ctx *Context) IsHelp() bool {
	return ctx.help
}

func (ctx *Context) Command() *Command {
	return ctx.command
}

func (ctx *Context) Completion() *Completion {
	return ctx.completion
}

func (ctx *Context) Flags() *FlagSet {
	return ctx.flags
}

func (ctx *Context) UnknownFlags() *FlagSet {
	return ctx.unknownFlags
}

//
// Before go into the sub command, we need traverse flags and merge with parent
func (ctx *Context) EnterCommand(cmd *Command) {
	ctx.command = cmd
	if !cmd.EnableUnknownFlag {
		ctx.unknownFlags = nil
	} else if ctx.unknownFlags == nil {
		ctx.unknownFlags = NewFlagSet()
	}

	ctx.flags = cmd.flags.mergeWith(ctx.flags, func(f *Flag) bool {
		return f.Persistent
	})
	ctx.flags.Add(HelpFlag)
}

func (ctx *Context) CheckFlags() error {
	for _, f := range ctx.flags.Flags() {
		if f.Required && !f.IsAssigned() {
			//if !f.useDefaultValue() {
			return fmt.Errorf("missing flag --%s", f.Name)
			//}
		}
		if err := f.checkFields(); err != nil {
			return err
		}

	}
	return nil
}

func (ctx *Context) detectFlag(name string) (*Flag, error) {
	flag := ctx.flags.Get(name)

	if flag != nil {
		return flag, nil
	} else if ctx.unknownFlags != nil {
		return ctx.unknownFlags.AddByName(name)
	} else {
		return nil, NewInvalidFlagError(name, ctx)
	}
}

func (ctx *Context) detectFlagByShorthand(ch rune) (*Flag, error) {
	flag := ctx.flags.GetByShorthand(ch)
	if flag != nil {
		return flag, nil
	} else {
		return nil, fmt.Errorf("unknown flag -%s", string(ch))
	}
}
<commit_msg>fix bugs with Flag.Fields checking<commit_after>/*
 * Copyright (C) 2017-2018 Alibaba Group Holding Limited
 */
package cli

import (
	"fmt"
	"github.com/aliyun/aliyun-cli/i18n"
)

//
// default help flag
var HelpFlag = &Flag{
	Name:         "help",
	Short:        i18n.T("print help", "打印帮助信息"),
	AssignedMode: AssignedNone,
}

//
// CLI Command Context
type Context struct {
	help         bool
	flags        *FlagSet
	unknownFlags *FlagSet
	command      *Command
	completion   *Completion
}

func NewCommandContext() *Context {
	return &Context{
		flags:        NewFlagSet(),
		unknownFlags: nil,
	}
}

func (ctx *Context) IsHelp() bool {
	return ctx.help
}

func (ctx *Context) Command() *Command {
	return ctx.command
}

func (ctx *Context) Completion() *Completion {
	return ctx.completion
}

func (ctx *Context) Flags() *FlagSet {
	return ctx.flags
}

func (ctx *Context) UnknownFlags() *FlagSet {
	return ctx.unknownFlags
}

//
// Before go into the sub command, we need traverse flags and merge with parent
func (ctx *Context) EnterCommand(cmd *Command) {
	ctx.command = cmd
	if !cmd.EnableUnknownFlag {
		ctx.unknownFlags = nil
	} else if ctx.unknownFlags == nil {
		ctx.unknownFlags = NewFlagSet()
	}

	ctx.flags = cmd.flags.mergeWith(ctx.flags, func(f *Flag) bool {
		return f.Persistent
	})
	ctx.flags.Add(HelpFlag)
}

func (ctx *Context) CheckFlags() error {
	for _, f := range ctx.flags.Flags() {
		if !f.IsAssigned() {
			if f.Required {
				return fmt.Errorf("missing flag --%s", f.Name)
			}
		} else {
			if err := f.checkFields(); err != nil {
				return err
			}
		}
	}
	return nil
}

func (ctx *Context) detectFlag(name string) (*Flag, error) {
	flag := ctx.flags.Get(name)

	if flag != nil {
		return flag, nil
	} else if ctx.unknownFlags != nil {
		return ctx.unknownFlags.AddByName(name)
	} else {
		return nil, NewInvalidFlagError(name, ctx)
	}
}

func (ctx *Context) detectFlagByShorthand(ch rune) (*Flag, error) {
	flag := ctx.flags.GetByShorthand(ch)
	if flag != nil {
		return flag, nil
	} else {
		return nil, fmt.Errorf("unknown flag -%s", string(ch))
	}
}
<|endoftext|>


<commit_before>package PoolAndAgent

import (
	"database/sql"
	"fmt"
	"strconv"
	"time"

	"github.com/3zheng/railgun/DialManager"
	ListenManager "github.com/3zheng/railgun/TcpListenManager"
	bs_proto "github.com/3zheng/railgun/protodefine"
	bs_types "github.com/3zheng/railgun/protodefine/mytype"
	bs_router "github.com/3zheng/railgun/protodefine/router"
	bs_tcp "github.com/3zheng/railgun/protodefine/tcpnet"

	_ "github.com/go-sql-driver/mysql"
	proto "github.com/golang/protobuf/proto"
)

type NetAgent struct {
	IpAddress string
}

var netAgentList []NetAgent

func (*NetAgent) SendMsg(req *bs_tcp.TCPTransferMsg) {
	var connId uint64 = req.Base.ConnId
	sess := ListenManager.GetSessionByConnId(connId)
	if sess == nil { //判断是否为空
		return
	}
	//往MsgWriteCh里写req
	select {
	case sess.MsgWriteCh <- req:
	case <-time.After(5 * time.Second):
		fmt.Println("严重事件，sess.MsgWriteCh已经阻塞超时5秒, connId =", connId)
	}
}

//创建一个NetAgent,ipAdd是个带端口的ip地址，如果是监听一个端口使用0.0.0.0:port
//NetAgent是TcpManager的映射，在CreateNetAgent并不和TcpManager联系起来，而是在SingleMsgPool的InitAndRun来创建TcpManager的TCP端口监听
func CreateNetAgent(ipAdd string) *NetAgent {
	for _, v := range netAgentList {
		if v.IpAddress == ipAdd {
			fmt.Println("已存在相同IP地址和端口")
			return nil
		}
	}
	agent := new(NetAgent)
	agent.IpAddress = ipAdd
	netAgentList = append(netAgentList, agent)
	return agent
}

type RouterAgent struct {
	IpAddress   string
	DialSession *DialManager.ConnectionSession
}

func CreateRouterAgent(ipAdd string) *RouterAgent {
	agent := new(RouterAgent)
	agent.IpAddress = ipAdd
	return agent
}

//启动运行RouterAgent
func (this *RouterAgent) RunRouterAgent(RouterToLogicChannel chan proto.Message, myAppId uint32, myAppType uint32) {
	fmt.Println("RunRouterAgent() myAppId=", myAppId, "myAppType=", myAppType)
	isClosed := false
	this.DialSession = DialManager.CreateClient(this.IpAddress, RouterToLogicChannel)
	for {
		//在这个router session关闭后，需要不停的尝试重新连接
		for isClosed || this.DialSession == nil {
			select {
			case <-time.After(700 * time.Millisecond): //每700毫秒尝试重新连接一次
			}
			this.DialSession = DialManager.CreateClient(this.IpAddress, RouterToLogicChannel)
			if this.DialSession != nil {
				//重连成功了
				fmt.Println("router Agent重连成功了")
				//这里用break是不能跳出for循环的，只能跳出当前的select，所以只能isClosed = false
				isClosed = false
			}
		}
		//连上router后向router发送注册报文
		req := new(bs_router.RegisterAppReq)
		bs_proto.SetBaseKindAndSubId(req)
		req.AppType = myAppType
		req.AppId = myAppId
		buff, err := proto.Marshal(req)
		if err != nil {
			fmt.Println("序列化RegisterAppReq出错")
		}
		tcpMsg := new(bs_tcp.TCPTransferMsg)
		bs_proto.SetBaseKindAndSubId(tcpMsg)
		tcpMsg.Data = buff
		tcpMsg.DataKindId = uint32(bs_types.CMDKindId_IDKindRouter)
		tcpMsg.DataSubId = uint32(bs_router.CMDID_Router_IDRegisterAppReq)
		this.SendMsg(tcpMsg)
		fmt.Println("向router app发送了注册请求")
		//阻塞在这里直到session关闭
		select {
		case v := <-this.DialSession.Quit:
			if v {
				fmt.Println("router agent的连接已关闭")
				isClosed = true
				this.DialSession = nil
			}
		}
	}
}

func (this *RouterAgent) SendMsg(req *bs_tcp.TCPTransferMsg) {
	var connId uint64 = req.Base.ConnId
	if this.DialSession == nil { //判断是否为空
		return
	}
	//往MsgWriteCh里写req
	select {
	case this.DialSession.MsgWriteCh <- req:
	case <-time.After(5 * time.Second):
		fmt.Println("严重事件，DialSession.MsgWriteCh已经阻塞超时5秒, connId =", connId)
	}

}

type DBReadInfo struct {
	TheRows    *sql.Rows      //如果是读操作，这里是返回的数据集
	TheColumns map[string]int //返回的列名,key是列名，value是表示这个是第几列，在返回的时候好查找
	ArrValues  [][]string     //返回的结果值
	RowNum     int            //行数，其实就是len(ArrValues)
}

func (this *DBReadInfo) Clear() {
	this.TheRows = nil
	this.TheColumns = nil
	this.ArrValues = nil
	this.RowNum = 0
}

type DBWriteInfo struct {
	LastId       int64 //自增列的当前值
	AffectedRows int64 //写操作影响的行数
}

//数据库的映射
type CADODatabase struct {
	DBSourceString string      //mysql的数据库连接串	格式"root:123456@tcp(localhost:3306)/sns?charset=utf8"
	Err            error       //错误
	TheDB          *sql.DB     //数据库对象
	ReadInfo       DBReadInfo  //如果是select的读操作，相关返回结果存这里
	WriteInfo      DBWriteInfo //如果是update或insert的写操作，相关返回结果存这里
}

func CreateADODatabase(DBSourceString string) *CADODatabase {
	DBAgent := new(CADODatabase)
	DBAgent.DBSourceString = DBSourceString
	return DBAgent
}

func (this *CADODatabase) InitDB() {
	db, err := sql.Open("mysql", this.DBSourceString)
	if err != nil {
		fmt.Println("初始化失败，error=", err.Error())
		panic(err.Error())
		this.Err = err
	}
	fmt.Println("数据库连接初始化成功")
	this.TheDB = db
}

//从数据库里读数据
func (this *CADODatabase) ReadFromDB(sqlExpress string) {
	//在读数据前先clear DBReadInfo
	this.ReadInfo.Clear()
	rows, err := this.TheDB.Query(sqlExpress)
	if err != nil {
		fmt.Println("读数据库出错,sql语句=", sqlExpress, "error=", err.Error())
		this.Err = err
		return
	}
	this.ReadInfo.TheRows = rows
	columns, err := rows.Columns()
	if err != nil {
		fmt.Println("取列名出错,sql语句=", sqlExpress, "error=", err.Error())
		this.Err = err
		return
	}

	this.ReadInfo.TheColumns = make(map[string]int)
	//记录列名及其对应的数组下标
	for i, v := range columns {
		this.ReadInfo.TheColumns[v] = i
	}
	colNum := len(columns) //列的数量
	// Make a slice for the values
	values := make([]sql.RawBytes, colNum)
	// rows.Scan wants '[]interface{}' as an argument, so we must copy the
	// references into such a slice
	// See http://code.google.com/p/go-wiki/wiki/InterfaceSlice for details
	scanArgs := make([]interface{}, colNum)
	for i := range values {
		scanArgs[i] = &values[i]
	}

	// Fetch rows
	for rows.Next() {
		// get RawBytes from data
		err = rows.Scan(scanArgs...)
		if err != nil {
			fmt.Println("从rows里取数据出错,sql语句=", sqlExpress, "error=", err.Error())
			this.Err = err
			return
		}

		// Now do something with the data.
		// Here we just print each column as a string.
		var value string
		rowValue := make([]string, colNum)
		for i, col := range values {
			// Here we can check if the value is nil (NULL value)
			if col == nil {
				value = "NULL"
			} else {
				value = string(col)
			}
			rowValue[i] = value
		}
		//把取出来的值放进二维数组中
		this.ReadInfo.ArrValues = append(this.ReadInfo.ArrValues, rowValue)
	}
	this.ReadInfo.RowNum = len(this.ReadInfo.ArrValues)
}

//向数据库写数据
func (this *CADODatabase) WriteToDB(sqlExpress string) {
	result, err := this.TheDB.Exec(sqlExpress)
	if err != nil {
		fmt.Println("写数据库出错,sql语句=", sqlExpress, "error=", err.Error())
		this.Err = err
		return
	}
	this.WriteInfo.LastId, _ = result.LastInsertId()
	this.WriteInfo.AffectedRows, _ = result.RowsAffected()
}

//在ReadFromDB从结果集里获取数据,结果放入value中，所以value要传地址。rowId从0开始为第一行
//成功返回true，失败返回false
func (this *CADODatabase) GetValueByRowIdAndColName(rowId int, colName string, value interface{}) bool {
	//先根据列名获得列的数组下标
	colId, ok := this.ReadInfo.TheColumns[colName]
	if !ok {
		fmt.Println("不存在的列名,colName=", colName)
		return false
	}
	if rowId >= this.ReadInfo.RowNum {
		fmt.Println("行号越界,rowId=", rowId, ",RowNum=", this.ReadInfo.RowNum)
		return false
	}
	//进行数据转换
	strValue := this.ReadInfo.ArrValues[rowId][colId]
	var err error
	switch data := value.(type) {
	case *int:
		*data, err = strconv.Atoi(strValue)
	case *int32:
		var v int64
		v, err = strconv.ParseInt(strValue, 10, 32)
		*data = int32(v)
	case *int64:
		*data, err = strconv.ParseInt(strValue, 10, 64)
	case *uint:
		var v uint64
		v, err = strconv.ParseUint(strValue, 10, 64)
		*data = uint(v)
	case *uint32:
		var v uint64
		v, err = strconv.ParseUint(strValue, 10, 32)
		*data = uint32(v)
	case *uint64:
		*data, err = strconv.ParseUint(strValue, 10, 64)
	case *float32:
		var v float64
		v, err = strconv.ParseFloat(strValue, 32)
		*data = float32(v)
	case *float64:
		*data, err = strconv.ParseFloat(strValue, 64)
	case *string:
		*data = strValue
	}
	if err != nil {
		fmt.Println("数据转换出错")
		return false
	}

	return true
}

type ILogicProcess interface {
	Init(myPool *SingleMsgPool) bool                       //初始化
	ProcessReq(req proto.Message, pDatabase *CADODatabase) //处理来自所绑定的SingleMsgPool发送的报文
	OnPulse(ms uint64)                                     //定时函数，每隔200ms调用一次
}
<commit_msg>避免重复创建NetAgent<commit_after>package PoolAndAgent

import (
	"database/sql"
	"fmt"
	"strconv"
	"time"

	"github.com/3zheng/railgun/DialManager"
	ListenManager "github.com/3zheng/railgun/TcpListenManager"
	bs_proto "github.com/3zheng/railgun/protodefine"
	bs_types "github.com/3zheng/railgun/protodefine/mytype"
	bs_router "github.com/3zheng/railgun/protodefine/router"
	bs_tcp "github.com/3zheng/railgun/protodefine/tcpnet"

	_ "github.com/go-sql-driver/mysql"
	proto "github.com/golang/protobuf/proto"
)

type NetAgent struct {
	IpAddress string
}

var netAgentList []NetAgent

func (*NetAgent) SendMsg(req *bs_tcp.TCPTransferMsg) {
	var connId uint64 = req.Base.ConnId
	sess := ListenManager.GetSessionByConnId(connId)
	if sess == nil { //判断是否为空
		return
	}
	//往MsgWriteCh里写req
	select {
	case sess.MsgWriteCh <- req:
	case <-time.After(5 * time.Second):
		fmt.Println("严重事件，sess.MsgWriteCh已经阻塞超时5秒, connId =", connId)
	}
}

//创建一个NetAgent,ipAdd是个带端口的ip地址，如果是监听一个端口使用0.0.0.0:port
//NetAgent是TcpManager的映射，在CreateNetAgent并不和TcpManager联系起来，而是在SingleMsgPool的InitAndRun来创建TcpManager的TCP端口监听
func CreateNetAgent(ipAdd string) *NetAgent {
	for _, v := range netAgentList {
		if v.IpAddress == ipAdd {
			fmt.Println("已存在相同IP地址和端口,返回空指针")
			return nil
		}
	}
	agent := new(NetAgent)
	agent.IpAddress = ipAdd
	netAgentList = append(netAgentList, agent)
	return agent
}

type RouterAgent struct {
	IpAddress   string
	DialSession *DialManager.ConnectionSession
}

func CreateRouterAgent(ipAdd string) *RouterAgent {
	agent := new(RouterAgent)
	agent.IpAddress = ipAdd
	return agent
}

//启动运行RouterAgent
func (this *RouterAgent) RunRouterAgent(RouterToLogicChannel chan proto.Message, myAppId uint32, myAppType uint32) {
	fmt.Println("RunRouterAgent() myAppId=", myAppId, "myAppType=", myAppType)
	isClosed := false
	this.DialSession = DialManager.CreateClient(this.IpAddress, RouterToLogicChannel)
	for {
		//在这个router session关闭后，需要不停的尝试重新连接
		for isClosed || this.DialSession == nil {
			select {
			case <-time.After(700 * time.Millisecond): //每700毫秒尝试重新连接一次
			}
			this.DialSession = DialManager.CreateClient(this.IpAddress, RouterToLogicChannel)
			if this.DialSession != nil {
				//重连成功了
				fmt.Println("router Agent重连成功了")
				//这里用break是不能跳出for循环的，只能跳出当前的select，所以只能isClosed = false
				isClosed = false
			}
		}
		//连上router后向router发送注册报文
		req := new(bs_router.RegisterAppReq)
		bs_proto.SetBaseKindAndSubId(req)
		req.AppType = myAppType
		req.AppId = myAppId
		buff, err := proto.Marshal(req)
		if err != nil {
			fmt.Println("序列化RegisterAppReq出错")
		}
		tcpMsg := new(bs_tcp.TCPTransferMsg)
		bs_proto.SetBaseKindAndSubId(tcpMsg)
		tcpMsg.Data = buff
		tcpMsg.DataKindId = uint32(bs_types.CMDKindId_IDKindRouter)
		tcpMsg.DataSubId = uint32(bs_router.CMDID_Router_IDRegisterAppReq)
		this.SendMsg(tcpMsg)
		fmt.Println("向router app发送了注册请求")
		//阻塞在这里直到session关闭
		select {
		case v := <-this.DialSession.Quit:
			if v {
				fmt.Println("router agent的连接已关闭")
				isClosed = true
				this.DialSession = nil
			}
		}
	}
}

func (this *RouterAgent) SendMsg(req *bs_tcp.TCPTransferMsg) {
	var connId uint64 = req.Base.ConnId
	if this.DialSession == nil { //判断是否为空
		return
	}
	//往MsgWriteCh里写req
	select {
	case this.DialSession.MsgWriteCh <- req:
	case <-time.After(5 * time.Second):
		fmt.Println("严重事件，DialSession.MsgWriteCh已经阻塞超时5秒, connId =", connId)
	}

}

type DBReadInfo struct {
	TheRows    *sql.Rows      //如果是读操作，这里是返回的数据集
	TheColumns map[string]int //返回的列名,key是列名，value是表示这个是第几列，在返回的时候好查找
	ArrValues  [][]string     //返回的结果值
	RowNum     int            //行数，其实就是len(ArrValues)
}

func (this *DBReadInfo) Clear() {
	this.TheRows = nil
	this.TheColumns = nil
	this.ArrValues = nil
	this.RowNum = 0
}

type DBWriteInfo struct {
	LastId       int64 //自增列的当前值
	AffectedRows int64 //写操作影响的行数
}

//数据库的映射
type CADODatabase struct {
	DBSourceString string      //mysql的数据库连接串	格式"root:123456@tcp(localhost:3306)/sns?charset=utf8"
	Err            error       //错误
	TheDB          *sql.DB     //数据库对象
	ReadInfo       DBReadInfo  //如果是select的读操作，相关返回结果存这里
	WriteInfo      DBWriteInfo //如果是update或insert的写操作，相关返回结果存这里
}

func CreateADODatabase(DBSourceString string) *CADODatabase {
	DBAgent := new(CADODatabase)
	DBAgent.DBSourceString = DBSourceString
	return DBAgent
}

func (this *CADODatabase) InitDB() {
	db, err := sql.Open("mysql", this.DBSourceString)
	if err != nil {
		fmt.Println("初始化失败，error=", err.Error())
		panic(err.Error())
		this.Err = err
	}
	fmt.Println("数据库连接初始化成功")
	this.TheDB = db
}

//从数据库里读数据
func (this *CADODatabase) ReadFromDB(sqlExpress string) {
	//在读数据前先clear DBReadInfo
	this.ReadInfo.Clear()
	rows, err := this.TheDB.Query(sqlExpress)
	if err != nil {
		fmt.Println("读数据库出错,sql语句=", sqlExpress, "error=", err.Error())
		this.Err = err
		return
	}
	this.ReadInfo.TheRows = rows
	columns, err := rows.Columns()
	if err != nil {
		fmt.Println("取列名出错,sql语句=", sqlExpress, "error=", err.Error())
		this.Err = err
		return
	}

	this.ReadInfo.TheColumns = make(map[string]int)
	//记录列名及其对应的数组下标
	for i, v := range columns {
		this.ReadInfo.TheColumns[v] = i
	}
	colNum := len(columns) //列的数量
	// Make a slice for the values
	values := make([]sql.RawBytes, colNum)
	// rows.Scan wants '[]interface{}' as an argument, so we must copy the
	// references into such a slice
	// See http://code.google.com/p/go-wiki/wiki/InterfaceSlice for details
	scanArgs := make([]interface{}, colNum)
	for i := range values {
		scanArgs[i] = &values[i]
	}

	// Fetch rows
	for rows.Next() {
		// get RawBytes from data
		err = rows.Scan(scanArgs...)
		if err != nil {
			fmt.Println("从rows里取数据出错,sql语句=", sqlExpress, "error=", err.Error())
			this.Err = err
			return
		}

		// Now do something with the data.
		// Here we just print each column as a string.
		var value string
		rowValue := make([]string, colNum)
		for i, col := range values {
			// Here we can check if the value is nil (NULL value)
			if col == nil {
				value = "NULL"
			} else {
				value = string(col)
			}
			rowValue[i] = value
		}
		//把取出来的值放进二维数组中
		this.ReadInfo.ArrValues = append(this.ReadInfo.ArrValues, rowValue)
	}
	this.ReadInfo.RowNum = len(this.ReadInfo.ArrValues)
}

//向数据库写数据
func (this *CADODatabase) WriteToDB(sqlExpress string) {
	result, err := this.TheDB.Exec(sqlExpress)
	if err != nil {
		fmt.Println("写数据库出错,sql语句=", sqlExpress, "error=", err.Error())
		this.Err = err
		return
	}
	this.WriteInfo.LastId, _ = result.LastInsertId()
	this.WriteInfo.AffectedRows, _ = result.RowsAffected()
}

//在ReadFromDB从结果集里获取数据,结果放入value中，所以value要传地址。rowId从0开始为第一行
//成功返回true，失败返回false
func (this *CADODatabase) GetValueByRowIdAndColName(rowId int, colName string, value interface{}) bool {
	//先根据列名获得列的数组下标
	colId, ok := this.ReadInfo.TheColumns[colName]
	if !ok {
		fmt.Println("不存在的列名,colName=", colName)
		return false
	}
	if rowId >= this.ReadInfo.RowNum {
		fmt.Println("行号越界,rowId=", rowId, ",RowNum=", this.ReadInfo.RowNum)
		return false
	}
	//进行数据转换
	strValue := this.ReadInfo.ArrValues[rowId][colId]
	var err error
	switch data := value.(type) {
	case *int:
		*data, err = strconv.Atoi(strValue)
	case *int32:
		var v int64
		v, err = strconv.ParseInt(strValue, 10, 32)
		*data = int32(v)
	case *int64:
		*data, err = strconv.ParseInt(strValue, 10, 64)
	case *uint:
		var v uint64
		v, err = strconv.ParseUint(strValue, 10, 64)
		*data = uint(v)
	case *uint32:
		var v uint64
		v, err = strconv.ParseUint(strValue, 10, 32)
		*data = uint32(v)
	case *uint64:
		*data, err = strconv.ParseUint(strValue, 10, 64)
	case *float32:
		var v float64
		v, err = strconv.ParseFloat(strValue, 32)
		*data = float32(v)
	case *float64:
		*data, err = strconv.ParseFloat(strValue, 64)
	case *string:
		*data = strValue
	}
	if err != nil {
		fmt.Println("数据转换出错")
		return false
	}

	return true
}

type ILogicProcess interface {
	Init(myPool *SingleMsgPool) bool                       //初始化
	ProcessReq(req proto.Message, pDatabase *CADODatabase) //处理来自所绑定的SingleMsgPool发送的报文
	OnPulse(ms uint64)                                     //定时函数，每隔200ms调用一次
}
<|endoftext|>


<commit_before>package tests

import (
	"errors"
	"fmt"
	"log"
	"path/filepath"
	"testing"

	"golang.org/x/net/publicsuffix"

	"github.com/BurntSushi/toml"
	"github.com/amadeovezz/gobro/db"
	"github.com/amadeovezz/gobro/parse"
)

var conf Config

type Config struct {
	Logs   logs
	Parser map[string]parser
	DB     database `toml:"database"`
}

type parser struct {
	Fields []string
}

type database struct {
	Username     string
	Password     string
	IP           string
	Port         string
	DatabaseName string
}

type logs struct {
	PathToConn string
	PathToDns  string
}

func (c *Config) SetupConfig(path string) {
	filename, _ := filepath.Abs(path)
	if _, err := toml.DecodeFile(filename, c); err != nil {
		log.Fatal("setupConfig() could not decode toml, err: ", err)
	}
}

func TestMain(m *testing.M) {

	conf.SetupConfig("config.toml")

	err := db.InitDB(
		conf.DB.Username,
		conf.DB.Password,
		conf.DB.IP,
		conf.DB.Port,
		conf.DB.DatabaseName,
	)

	if err != nil {
		fmt.Println(err)
		return
	}

	m.Run()

}

func TestParseConn(t *testing.T) {

	// Create a new parser with specific fields, and raw entries
	parser, err := parse.NewParser(conf.Logs.PathToConn, false)
	if err != nil {
		t.Fatal(err)
	}

	// Grab the specific fields to parse from config
	parser.SetFields(conf.Parser["conn"].Fields)

	// How many rows do you want to buffer
	parser.CreateBuffer(100)

	go parser.BufferRow()

	err = db.InsertBatch(parser.Row, "conn", len(parser.Fields()))
	if err != nil {
		t.Fatal(err)
	}

	log.Fatal("reached parseconn")
}

func DnsParse(fields, row []string) ([]string, error) {

	var newRow []string
	newRow = row

	for i, field := range fields {

		if field == "query" {
			if row[i] == "-" {
				// Maybe change this to bool
				return nil, errors.New("No query information provided")
			}
			secondLevelDomain, err := publicsuffix.EffectiveTLDPlusOne(newRow[i])
			if err == nil {
				newRow[i] = secondLevelDomain
			}
		}

	}

	return newRow, nil

}

func TestParseDns(t *testing.T) {

	// Create a new parser with specific field and augemented entries
	parser, err := parse.NewParser(conf.Logs.PathToDns, false)
	if err != nil {
		t.Fatal(err)
	}

	// Grab the specific fields to parse from config
	parser.SetFields(conf.Parser["dns"].Fields)

	// How many rows do you want to buffer
	parser.CreateBuffer(100)

	// Strip out uncessary domain information
	go parser.BufferRow(DnsParse)

	// Insert into db
	err = db.InsertBatch(parser.Row, "dns", len(parser.Fields()))
	if err != nil {
		t.Fatal(err)
	}

}

func BenchmarkParsingAndInsertingDb(b *testing.B) {

	for i := 0; i < b.N; i++ {

		parser, err := parse.NewParser(conf.Logs.PathToDns, false)
		if err != nil {
			b.Fatal(err)
		}

		parser.SetFields(conf.Parser["dns"].Fields)

		parser.AutoCreateBuffer()
		parser.BufferRow()

		// Insert into db
		err = db.InsertBatch(parser.Row, "dns", len(parser.Fields()))
		if err != nil {
			b.Fatal(err)
		}

	}

}
<commit_msg>Removed fatal log<commit_after>package tests

import (
	"errors"
	"fmt"
	"log"
	"path/filepath"
	"testing"

	"golang.org/x/net/publicsuffix"

	"github.com/BurntSushi/toml"
	"github.com/amadeovezz/gobro/db"
	"github.com/amadeovezz/gobro/parse"
)

var conf Config

type Config struct {
	Logs   logs
	Parser map[string]parser
	DB     database `toml:"database"`
}

type parser struct {
	Fields []string
}

type database struct {
	Username     string
	Password     string
	IP           string
	Port         string
	DatabaseName string
}

type logs struct {
	PathToConn string
	PathToDns  string
}

func (c *Config) SetupConfig(path string) {
	filename, _ := filepath.Abs(path)
	if _, err := toml.DecodeFile(filename, c); err != nil {
		log.Fatal("setupConfig() could not decode toml, err: ", err)
	}
}

func TestMain(m *testing.M) {

	conf.SetupConfig("config.toml")

	err := db.InitDB(
		conf.DB.Username,
		conf.DB.Password,
		conf.DB.IP,
		conf.DB.Port,
		conf.DB.DatabaseName,
	)

	if err != nil {
		fmt.Println(err)
		return
	}

	m.Run()

}

func TestParseConn(t *testing.T) {

	// Create a new parser with specific fields, and raw entries
	parser, err := parse.NewParser(conf.Logs.PathToConn, false)
	if err != nil {
		t.Fatal(err)
	}

	// Grab the specific fields to parse from config
	parser.SetFields(conf.Parser["conn"].Fields)

	// How many rows do you want to buffer
	parser.CreateBuffer(100)

	go parser.BufferRow()

	err = db.InsertBatch(parser.Row, "conn", len(parser.Fields()))
	if err != nil {
		t.Fatal(err)
	}
}

func DnsParse(fields, row []string) ([]string, error) {

	var newRow []string
	newRow = row

	for i, field := range fields {

		if field == "query" {
			if row[i] == "-" {
				// Maybe change this to bool
				return nil, errors.New("No query information provided")
			}
			secondLevelDomain, err := publicsuffix.EffectiveTLDPlusOne(newRow[i])
			if err == nil {
				newRow[i] = secondLevelDomain
			}
		}

	}

	return newRow, nil

}

func TestParseDns(t *testing.T) {

	// Create a new parser with specific field and augemented entries
	parser, err := parse.NewParser(conf.Logs.PathToDns, false)
	if err != nil {
		t.Fatal(err)
	}

	// Grab the specific fields to parse from config
	parser.SetFields(conf.Parser["dns"].Fields)

	// How many rows do you want to buffer
	parser.CreateBuffer(100)

	// Strip out uncessary domain information
	go parser.BufferRow(DnsParse)

	// Insert into db
	err = db.InsertBatch(parser.Row, "dns", len(parser.Fields()))
	if err != nil {
		t.Fatal(err)
	}

}

func BenchmarkParsingAndInsertingDb(b *testing.B) {

	for i := 0; i < b.N; i++ {

		parser, err := parse.NewParser(conf.Logs.PathToDns, false)
		if err != nil {
			b.Fatal(err)
		}

		parser.SetFields(conf.Parser["dns"].Fields)

		parser.AutoCreateBuffer()
		parser.BufferRow()

		// Insert into db
		err = db.InsertBatch(parser.Row, "dns", len(parser.Fields()))
		if err != nil {
			b.Fatal(err)
		}

	}

}
<|endoftext|>


<commit_before>package client

import (
	"fmt"
	"strconv"
	"strings"
	"syscall"

	gflag "github.com/jessevdk/go-flags"
)

func (cli *HyperClient) HyperCmdKill(args ...string) error {
	var opts struct {
		Pod    bool   `short:"p" long:"pod" default:"false" default-mask:"-" description:"kill all containers in a pod"`
		Signal string `short:"s" long:"signal" value-name:"\"\"" description:"The signal to kill containers, default is 9"`
	}
	var parser = gflag.NewParser(&opts, gflag.Default|gflag.IgnoreUnknown)
	parser.Usage = "kill [OPTIONS] CONTAINER_ID|POD_ID\n\nSend kill signal to container or Pod"
	args, err := parser.ParseArgs(args)
	if err != nil {
		if !strings.Contains(err.Error(), "Usage") {
			return err
		} else {
			return nil
		}
	}

	if len(args) == 0 {
		return fmt.Errorf("\"kill\" requires a minimum of 1 argument, please provide VM ID.\n")
	}

	sig := 9
	if opts.Signal != "" {
		s, err := parseSignal(opts.Signal)
		if err != nil {
			return err
		}
		sig = int(s)
	}

	for i := range args {
		if opts.Pod {
			err := cli.client.KillPod(args[i], sig)
			fmt.Fprintf(cli.err, "failed to delete pod %s: %v", args[i], err)
		} else {
			err := cli.client.KillContainer(args[i], sig)
			fmt.Fprintf(cli.err, "failed to delete container %s: %v", args[i], err)
		}
	}

	return nil
}

func parseSignal(rawSignal string) (syscall.Signal, error) {
	s, err := strconv.Atoi(rawSignal)
	if err == nil {
		return syscall.Signal(s), nil
	}
	signal, ok := linuxSignalMap[strings.TrimPrefix(strings.ToUpper(rawSignal), "SIG")]
	if !ok {
		return -1, fmt.Errorf("unknown signal %q", rawSignal)
	}
	return signal, nil
}

var linuxSignalMap = map[string]syscall.Signal{
	"ABRT":   syscall.SIGABRT,
	"ALRM":   syscall.SIGALRM,
	"BUS":    syscall.SIGBUS,
	"CHLD":   syscall.SIGCHLD,
	"CLD":    syscall.SIGCLD,
	"CONT":   syscall.SIGCONT,
	"FPE":    syscall.SIGFPE,
	"HUP":    syscall.SIGHUP,
	"ILL":    syscall.SIGILL,
	"INT":    syscall.SIGINT,
	"IO":     syscall.SIGIO,
	"IOT":    syscall.SIGIOT,
	"KILL":   syscall.SIGKILL,
	"PIPE":   syscall.SIGPIPE,
	"POLL":   syscall.SIGPOLL,
	"PROF":   syscall.SIGPROF,
	"PWR":    syscall.SIGPWR,
	"QUIT":   syscall.SIGQUIT,
	"SEGV":   syscall.SIGSEGV,
	"STKFLT": syscall.SIGSTKFLT,
	"STOP":   syscall.SIGSTOP,
	"SYS":    syscall.SIGSYS,
	"TERM":   syscall.SIGTERM,
	"TRAP":   syscall.SIGTRAP,
	"TSTP":   syscall.SIGTSTP,
	"TTIN":   syscall.SIGTTIN,
	"TTOU":   syscall.SIGTTOU,
	"UNUSED": syscall.SIGUNUSED,
	"URG":    syscall.SIGURG,
	"USR1":   syscall.SIGUSR1,
	"USR2":   syscall.SIGUSR2,
	"VTALRM": syscall.SIGVTALRM,
	"WINCH":  syscall.SIGWINCH,
	"XCPU":   syscall.SIGXCPU,
	"XFSZ":   syscall.SIGXFSZ,
}
<commit_msg>fix output message of hyperctl kill<commit_after>package client

import (
	"fmt"
	"strconv"
	"strings"
	"syscall"

	gflag "github.com/jessevdk/go-flags"
)

func (cli *HyperClient) HyperCmdKill(args ...string) error {
	var opts struct {
		Pod    bool   `short:"p" long:"pod" default:"false" default-mask:"-" description:"kill all containers in a pod"`
		Signal string `short:"s" long:"signal" value-name:"\"\"" description:"The signal to kill containers, default is 9"`
	}
	var parser = gflag.NewParser(&opts, gflag.Default|gflag.IgnoreUnknown)
	parser.Usage = "kill [OPTIONS] CONTAINER_ID|POD_ID\n\nSend kill signal to container or Pod"
	args, err := parser.ParseArgs(args)
	if err != nil {
		if !strings.Contains(err.Error(), "Usage") {
			return err
		} else {
			return nil
		}
	}

	if len(args) == 0 {
		return fmt.Errorf("\"kill\" requires a minimum of 1 argument, please provide container ID.\n")
	}

	sig := 9
	if opts.Signal != "" {
		s, err := parseSignal(opts.Signal)
		if err != nil {
			return err
		}
		sig = int(s)
	}

	for i := range args {
		if opts.Pod {
			err = cli.client.KillPod(args[i], sig)
			if err != nil {
				fmt.Fprintf(cli.err, "failed to kill pod %s: %v", args[i], err)
			}
		} else {
			err = cli.client.KillContainer(args[i], sig)
			if err != nil {
				fmt.Fprintf(cli.err, "failed to kill container %s: %v", args[i], err)
			}
		}
	}

	return nil
}

func parseSignal(rawSignal string) (syscall.Signal, error) {
	s, err := strconv.Atoi(rawSignal)
	if err == nil {
		return syscall.Signal(s), nil
	}
	signal, ok := linuxSignalMap[strings.TrimPrefix(strings.ToUpper(rawSignal), "SIG")]
	if !ok {
		return -1, fmt.Errorf("unknown signal %q", rawSignal)
	}
	return signal, nil
}

var linuxSignalMap = map[string]syscall.Signal{
	"ABRT":   syscall.SIGABRT,
	"ALRM":   syscall.SIGALRM,
	"BUS":    syscall.SIGBUS,
	"CHLD":   syscall.SIGCHLD,
	"CLD":    syscall.SIGCLD,
	"CONT":   syscall.SIGCONT,
	"FPE":    syscall.SIGFPE,
	"HUP":    syscall.SIGHUP,
	"ILL":    syscall.SIGILL,
	"INT":    syscall.SIGINT,
	"IO":     syscall.SIGIO,
	"IOT":    syscall.SIGIOT,
	"KILL":   syscall.SIGKILL,
	"PIPE":   syscall.SIGPIPE,
	"POLL":   syscall.SIGPOLL,
	"PROF":   syscall.SIGPROF,
	"PWR":    syscall.SIGPWR,
	"QUIT":   syscall.SIGQUIT,
	"SEGV":   syscall.SIGSEGV,
	"STKFLT": syscall.SIGSTKFLT,
	"STOP":   syscall.SIGSTOP,
	"SYS":    syscall.SIGSYS,
	"TERM":   syscall.SIGTERM,
	"TRAP":   syscall.SIGTRAP,
	"TSTP":   syscall.SIGTSTP,
	"TTIN":   syscall.SIGTTIN,
	"TTOU":   syscall.SIGTTOU,
	"UNUSED": syscall.SIGUNUSED,
	"URG":    syscall.SIGURG,
	"USR1":   syscall.SIGUSR1,
	"USR2":   syscall.SIGUSR2,
	"VTALRM": syscall.SIGVTALRM,
	"WINCH":  syscall.SIGWINCH,
	"XCPU":   syscall.SIGXCPU,
	"XFSZ":   syscall.SIGXFSZ,
}
<|endoftext|>


<commit_before>// flashlight is a lightweight chained proxy that can run in client or server mode.
package main

import (
	"flag"
	"math/rand"
	"os"
	"runtime"
	"time"

	"github.com/getlantern/fronted"
	"github.com/getlantern/golog"
	"github.com/getlantern/profiling"

	"github.com/getlantern/flashlight/client"
	"github.com/getlantern/flashlight/config"
	"github.com/getlantern/flashlight/http"
	"github.com/getlantern/flashlight/server"
	"github.com/getlantern/flashlight/statreporter"
	"github.com/getlantern/flashlight/statserver"
	"github.com/getlantern/whitelist"
)

const (
	// Exit Statuses
	ConfigError    = 1
	PortmapFailure = 50
)

var (
	log       = golog.LoggerFor("flashlight")
	version   string
	buildDate string

	// Command-line Flags
	help      = flag.Bool("help", false, "Get usage help")
	parentPID = flag.Int("parentpid", 0, "the parent process's PID, used on Windows for killing flashlight when the parent disappears")

	configUpdates = make(chan *config.Config)
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

func main() {
	displayVersion()

	flag.Parse()
	configUpdates = make(chan *config.Config)
	cfg, err := config.Start(func(updated *config.Config) {
		configUpdates <- updated
	})
	if err != nil {
		log.Fatalf("Unable to start configuration: %s", err)
	}
	if *help || cfg.Addr == "" || (cfg.Role != "server" && cfg.Role != "client") {
		flag.Usage()
		os.Exit(ConfigError)
	}

	finishProfiling := profiling.Start(cfg.CpuProfile, cfg.MemProfile)
	defer finishProfiling()

	// Configure stats initially
	configureStats(cfg, true)

	log.Debugf("Running proxy")
	if cfg.IsDownstream() {
		runClientProxy(cfg)
	} else {
		runServerProxy(cfg)
	}
}

func displayVersion() {
	if version == "" {
		version = "development"
	}
	if buildDate == "" {
		buildDate = "now"
	}
	log.Debugf("---- flashlight version %s (%s) ----", version, buildDate)
}

func configureStats(cfg *config.Config, failOnError bool) {
	err := statreporter.Configure(cfg.Stats)
	if err != nil {
		log.Error(err)
		if failOnError {
			flag.Usage()
			os.Exit(ConfigError)
		}
	}

	if cfg.StatsAddr != "" {
		statserver.Start(cfg.StatsAddr)
	} else {
		statserver.Stop()
	}
}

// Runs the client-side proxy
func runClientProxy(cfg *config.Config) {
	client := &client.Client{
		Addr:         cfg.Addr,
		ReadTimeout:  0, // don't timeout
		WriteTimeout: 0,
	}

	// Configure client initially
	client.Configure(cfg.Client)

	// this flag specifies the port to open the HTTP server
	// between the UI and
	// with a corresponding HTTP server at
	// the following address
	if cfg.Client.UiAddr != "" {
		wlChan := make(chan *whitelist.Config)
		go func() {
			err := http.UiHttpServer(cfg.Client, configUpdates, wlChan)
			if err != nil {
				log.Fatalf("Unable to start UI http server: %s", err)
			}
		}()

		go func() {
			for {
				// separate channel for synchronizing whitelist updates
				wl := <-wlChan
				cfg.Client.Whitelist = wl
				err := config.Update(func(updated *config.Config) error {
					log.Debugf("Saving updated whitelist configuration")
					updated.Client.Whitelist = cfg.Client.Whitelist
					return nil
				})

				if err != nil {
					log.Errorf("Could not update config: %s", err)
				}
			}
		}()
	}

	// Continually poll for config updates and update client accordingly
	go func() {
		for {
			cfg := <-configUpdates
			configureStats(cfg, false)
			client.Configure(cfg.Client)
		}
	}()

	err := client.ListenAndServe()
	if err != nil {
		log.Fatalf("Unable to run client proxy: %s", err)
	}
}

// Runs the server-side proxy
func runServerProxy(cfg *config.Config) {
	useAllCores()

	srv := &server.Server{
		Addr:         cfg.Addr,
		Host:         cfg.Server.AdvertisedHost,
		ReadTimeout:  0, // don't timeout
		WriteTimeout: 0,
		CertContext: &fronted.CertContext{
			PKFile:         config.InConfigDir("proxypk.pem"),
			ServerCertFile: config.InConfigDir("servercert.pem"),
		},
		AllowedPorts: []int{80, 443, 8080, 8443, 5222},
	}

	srv.Configure(cfg.Server)

	// Continually poll for config updates and update server accordingly
	go func() {
		for {
			cfg := <-configUpdates
			configureStats(cfg, false)
			srv.Configure(cfg.Server)
		}
	}()

	err := srv.ListenAndServe()
	if err != nil {
		log.Fatalf("Unable to run server proxy: %s", err)
	}
}

func useAllCores() {
	numcores := runtime.NumCPU()
	log.Debugf("Using all %d cores on machine", numcores)
	runtime.GOMAXPROCS(numcores)
}
<commit_msg>fix UI capitalization<commit_after>// flashlight is a lightweight chained proxy that can run in client or server mode.
package main

import (
	"flag"
	"math/rand"
	"os"
	"runtime"
	"time"

	"github.com/getlantern/fronted"
	"github.com/getlantern/golog"
	"github.com/getlantern/profiling"

	"github.com/getlantern/flashlight/client"
	"github.com/getlantern/flashlight/config"
	"github.com/getlantern/flashlight/http"
	"github.com/getlantern/flashlight/server"
	"github.com/getlantern/flashlight/statreporter"
	"github.com/getlantern/flashlight/statserver"
	"github.com/getlantern/whitelist"
)

const (
	// Exit Statuses
	ConfigError    = 1
	PortmapFailure = 50
)

var (
	log       = golog.LoggerFor("flashlight")
	version   string
	buildDate string

	// Command-line Flags
	help      = flag.Bool("help", false, "Get usage help")
	parentPID = flag.Int("parentpid", 0, "the parent process's PID, used on Windows for killing flashlight when the parent disappears")

	configUpdates = make(chan *config.Config)
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

func main() {
	displayVersion()

	flag.Parse()
	configUpdates = make(chan *config.Config)
	cfg, err := config.Start(func(updated *config.Config) {
		configUpdates <- updated
	})
	if err != nil {
		log.Fatalf("Unable to start configuration: %s", err)
	}
	if *help || cfg.Addr == "" || (cfg.Role != "server" && cfg.Role != "client") {
		flag.Usage()
		os.Exit(ConfigError)
	}

	finishProfiling := profiling.Start(cfg.CpuProfile, cfg.MemProfile)
	defer finishProfiling()

	// Configure stats initially
	configureStats(cfg, true)

	log.Debugf("Running proxy")
	if cfg.IsDownstream() {
		runClientProxy(cfg)
	} else {
		runServerProxy(cfg)
	}
}

func displayVersion() {
	if version == "" {
		version = "development"
	}
	if buildDate == "" {
		buildDate = "now"
	}
	log.Debugf("---- flashlight version %s (%s) ----", version, buildDate)
}

func configureStats(cfg *config.Config, failOnError bool) {
	err := statreporter.Configure(cfg.Stats)
	if err != nil {
		log.Error(err)
		if failOnError {
			flag.Usage()
			os.Exit(ConfigError)
		}
	}

	if cfg.StatsAddr != "" {
		statserver.Start(cfg.StatsAddr)
	} else {
		statserver.Stop()
	}
}

// Runs the client-side proxy
func runClientProxy(cfg *config.Config) {
	client := &client.Client{
		Addr:         cfg.Addr,
		ReadTimeout:  0, // don't timeout
		WriteTimeout: 0,
	}

	// Configure client initially
	client.Configure(cfg.Client)

	// this flag specifies the port to open the HTTP server
	// between the UI and
	// with a corresponding HTTP server at
	// the following address
	if cfg.Client.UIAddr != "" {
		wlChan := make(chan *whitelist.Config)
		go func() {
			err := http.UiHttpServer(cfg.Client, configUpdates, wlChan)
			if err != nil {
				log.Fatalf("Unable to start UI http server: %s", err)
			}
		}()

		go func() {
			for {
				// separate channel for synchronizing whitelist updates
				wl := <-wlChan
				cfg.Client.Whitelist = wl
				err := config.Update(func(updated *config.Config) error {
					log.Debugf("Saving updated whitelist configuration")
					updated.Client.Whitelist = cfg.Client.Whitelist
					return nil
				})

				if err != nil {
					log.Errorf("Could not update config: %s", err)
				}
			}
		}()
	}

	// Continually poll for config updates and update client accordingly
	go func() {
		for {
			cfg := <-configUpdates
			configureStats(cfg, false)
			client.Configure(cfg.Client)
		}
	}()

	err := client.ListenAndServe()
	if err != nil {
		log.Fatalf("Unable to run client proxy: %s", err)
	}
}

// Runs the server-side proxy
func runServerProxy(cfg *config.Config) {
	useAllCores()

	srv := &server.Server{
		Addr:         cfg.Addr,
		Host:         cfg.Server.AdvertisedHost,
		ReadTimeout:  0, // don't timeout
		WriteTimeout: 0,
		CertContext: &fronted.CertContext{
			PKFile:         config.InConfigDir("proxypk.pem"),
			ServerCertFile: config.InConfigDir("servercert.pem"),
		},
		AllowedPorts: []int{80, 443, 8080, 8443, 5222},
	}

	srv.Configure(cfg.Server)

	// Continually poll for config updates and update server accordingly
	go func() {
		for {
			cfg := <-configUpdates
			configureStats(cfg, false)
			srv.Configure(cfg.Server)
		}
	}()

	err := srv.ListenAndServe()
	if err != nil {
		log.Fatalf("Unable to run server proxy: %s", err)
	}
}

func useAllCores() {
	numcores := runtime.NumCPU()
	log.Debugf("Using all %d cores on machine", numcores)
	runtime.GOMAXPROCS(numcores)
}
<|endoftext|>