<commit_before>package http

import (
	"context"
	"io"
	"io/ioutil"
	"net/http"

	"github.com/cloudevents/sdk-go/v2/binding"
	"github.com/cloudevents/sdk-go/v2/binding/format"
	"github.com/cloudevents/sdk-go/v2/binding/spec"
	"github.com/cloudevents/sdk-go/v2/types"
)

// Fill the provided httpRequest with the message m.
// Using context you can tweak the encoding processing (more details on binding.Write documentation).
func WriteRequest(ctx context.Context, m binding.Message, httpRequest *http.Request, transformers ...binding.TransformerFactory) error {
	structuredWriter := (*httpRequestWriter)(httpRequest)
	binaryWriter := (*httpRequestWriter)(httpRequest)

	_, err := binding.Write(
		ctx,
		m,
		structuredWriter,
		binaryWriter,
		transformers...,
	)
	return err
}

type httpRequestWriter http.Request

func (b *httpRequestWriter) SetStructuredEvent(ctx context.Context, format format.Format, event io.Reader) error {
	b.Header.Set(ContentType, format.MediaType())
	b.Body = ioutil.NopCloser(event)
	return nil
}

func (b *httpRequestWriter) Start(ctx context.Context) error {
	return nil
}

func (b *httpRequestWriter) End(ctx context.Context) error {
	return nil
}

func (b *httpRequestWriter) SetData(reader io.Reader) error {
	b.Body = ioutil.NopCloser(reader)
	return nil
}

func (b *httpRequestWriter) SetAttribute(attribute spec.Attribute, value interface{}) error {
	// Http headers, everything is a string!
	s, err := types.Format(value)
	if err != nil {
		return err
	}

	if attribute.Kind() == spec.DataContentType {
		b.Header.Add(ContentType, s)
	} else {
		b.Header.Add(prefix+attribute.Name(), s)
	}
	return nil
}

func (b *httpRequestWriter) SetExtension(name string, value interface{}) error {
	// Http headers, everything is a string!
	s, err := types.Format(value)
	if err != nil {
		return err
	}
	b.Header.Add(prefix+name, s)
	return nil
}

var _ binding.StructuredWriter = (*httpRequestWriter)(nil) // Test it conforms to the interface
var _ binding.BinaryWriter = (*httpRequestWriter)(nil)     // Test it conforms to the interface
<commit_msg>write out content length for http request (#405)<commit_after>package http

import (
	"bytes"
	"context"
	"io"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/cloudevents/sdk-go/v2/binding"
	"github.com/cloudevents/sdk-go/v2/binding/format"
	"github.com/cloudevents/sdk-go/v2/binding/spec"
	"github.com/cloudevents/sdk-go/v2/types"
)

// Fill the provided httpRequest with the message m.
// Using context you can tweak the encoding processing (more details on binding.Write documentation).
func WriteRequest(ctx context.Context, m binding.Message, httpRequest *http.Request, transformers ...binding.TransformerFactory) error {
	structuredWriter := (*httpRequestWriter)(httpRequest)
	binaryWriter := (*httpRequestWriter)(httpRequest)

	_, err := binding.Write(
		ctx,
		m,
		structuredWriter,
		binaryWriter,
		transformers...,
	)
	return err
}

type httpRequestWriter http.Request

func (b *httpRequestWriter) SetStructuredEvent(ctx context.Context, format format.Format, event io.Reader) error {
	b.Header.Set(ContentType, format.MediaType())
	return b.setBody(event)
}

func (b *httpRequestWriter) Start(ctx context.Context) error {
	return nil
}

func (b *httpRequestWriter) End(ctx context.Context) error {
	return nil
}

func (b *httpRequestWriter) SetData(data io.Reader) error {
	return b.setBody(data)
}

// setBody is a cherry-pick of the implementation in http.NewRequestWithContext
func (b *httpRequestWriter) setBody(body io.Reader) error {
	rc, ok := body.(io.ReadCloser)
	if !ok && body != nil {
		rc = ioutil.NopCloser(body)
	}
	b.Body = rc
	if body != nil {
		switch v := body.(type) {
		case *bytes.Buffer:
			b.ContentLength = int64(v.Len())
			buf := v.Bytes()
			b.GetBody = func() (io.ReadCloser, error) {
				r := bytes.NewReader(buf)
				return ioutil.NopCloser(r), nil
			}
		case *bytes.Reader:
			b.ContentLength = int64(v.Len())
			snapshot := *v
			b.GetBody = func() (io.ReadCloser, error) {
				r := snapshot
				return ioutil.NopCloser(&r), nil
			}
		case *strings.Reader:
			b.ContentLength = int64(v.Len())
			snapshot := *v
			b.GetBody = func() (io.ReadCloser, error) {
				r := snapshot
				return ioutil.NopCloser(&r), nil
			}
		default:
			// This is where we'd set it to -1 (at least
			// if body != NoBody) to mean unknown, but
			// that broke people during the Go 1.8 testing
			// period. People depend on it being 0 I
			// guess. Maybe retry later. See Issue 18117.
		}
		// For client requests, Request.ContentLength of 0
		// means either actually 0, or unknown. The only way
		// to explicitly say that the ContentLength is zero is
		// to set the Body to nil. But turns out too much code
		// depends on NewRequest returning a non-nil Body,
		// so we use a well-known ReadCloser variable instead
		// and have the http package also treat that sentinel
		// variable to mean explicitly zero.
		if b.GetBody != nil && b.ContentLength == 0 {
			b.Body = http.NoBody
			b.GetBody = func() (io.ReadCloser, error) { return http.NoBody, nil }
		}
	}
	return nil
}

func (b *httpRequestWriter) SetAttribute(attribute spec.Attribute, value interface{}) error {
	// Http headers, everything is a string!
	s, err := types.Format(value)
	if err != nil {
		return err
	}

	if attribute.Kind() == spec.DataContentType {
		b.Header.Add(ContentType, s)
	} else {
		b.Header.Add(prefix+attribute.Name(), s)
	}
	return nil
}

func (b *httpRequestWriter) SetExtension(name string, value interface{}) error {
	// Http headers, everything is a string!
	s, err := types.Format(value)
	if err != nil {
		return err
	}
	b.Header.Add(prefix+name, s)
	return nil
}

var _ binding.StructuredWriter = (*httpRequestWriter)(nil) // Test it conforms to the interface
var _ binding.BinaryWriter = (*httpRequestWriter)(nil)     // Test it conforms to the interface
<|endoftext|>


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

package s3

import (
	"fmt"
	"github.com/jacobsa/aws"
	"github.com/jacobsa/aws/s3/auth"
	"github.com/jacobsa/aws/s3/http"
	"github.com/jacobsa/aws/s3/time"
	"net/url"
	sys_time "time"
	"unicode/utf8"
)

// Bucket represents an S3 bucket, which is a collection of objects keyed on
// Unicode strings. The UTF-8 encoding of a key must be no more than 1024 bytes
// long.
//
// See here for more info:
//
//     http://goo.gl/Nd63t
//
type Bucket interface {
	// Retrieve data for the object with the given key.
	GetObject(key string) (data []byte, err error)

	// Store the supplied data with the given key, overwriting any previous
	// version. The object is created with the default ACL of "private".
	StoreObject(key string, data []byte) error
}

// OpenBucket returns a Bucket tied to a given name in whe given region. You
// must have previously created the bucket in the region, and the supplied
// access key must have access to it.
//
// To easily create a bucket, use the AWS Console:
//
//     http://aws.amazon.com/console/
//
func OpenBucket(name string, region Region, key aws.AccessKey) (Bucket, error) {
	// Create a connection to the given region's endpoint.
	endpoint := &url.URL{Scheme: "https", Host: string(region)}
	httpConn, err := http.NewConn(endpoint)
	if err != nil {
		return nil, fmt.Errorf("http.NewConn: %v", err)
	}

	// Create an appropriate request signer.
	signer, err := auth.NewSigner(&key)
	if err != nil {
		return nil, fmt.Errorf("auth.NewSigner: %v", err)
	}

	return openBucket(name, httpConn, signer, time.RealClock())
}

// A version of OpenBucket with the ability to inject dependencies, for
// testability.
func openBucket(
	name string,
	httpConn http.Conn,
	signer auth.Signer,
	clock time.Clock) (Bucket, error) {
	return &bucket{name, httpConn, signer, clock}, nil
}

type bucket struct {
	name     string
	httpConn http.Conn
	signer   auth.Signer
	clock    time.Clock
}

func (b *bucket) GetObject(key string) (data []byte, err error) {
	return nil, fmt.Errorf("TODO: Implement bucket.GetObject.")
}

func (b *bucket) StoreObject(key string, data []byte) error {
	// Validate the key.
	if len(key) > 1024 {
		return fmt.Errorf("Keys may be no longer than 1024 bytes.")
	}

	if !utf8.ValidString(key) {
		return fmt.Errorf("Keys must be valid UTF-8.")
	}

	// Build an appropriate HTTP request.
	httpReq := &http.Request{
		Verb: "PUT",
		Path: fmt.Sprintf("/%s/%s", b.name, key),
		Body: data,
		Headers: map[string]string{
			"Date": b.clock.Now().UTC().Format(sys_time.RFC1123),
		},
	}

	// Sign the request.
	if err := b.signer.Sign(httpReq); err != nil {
		return fmt.Errorf("Sign: %v", err)
	}

	// Send the request.
	httpResp, err := b.httpConn.SendRequest(httpReq)
	if err != nil {
		return fmt.Errorf("SendRequest: %v", err)
	}

	// Check the response.
	if httpResp.StatusCode != 200 {
		return fmt.Errorf("Error from server: %d %s", httpResp.StatusCode, httpResp.Body)
	}

	return nil
}
<commit_msg>Added a call-out to the reference.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package s3

import (
	"fmt"
	"github.com/jacobsa/aws"
	"github.com/jacobsa/aws/s3/auth"
	"github.com/jacobsa/aws/s3/http"
	"github.com/jacobsa/aws/s3/time"
	"net/url"
	sys_time "time"
	"unicode/utf8"
)

// Bucket represents an S3 bucket, which is a collection of objects keyed on
// Unicode strings. The UTF-8 encoding of a key must be no more than 1024 bytes
// long.
//
// See here for more info:
//
//     http://goo.gl/Nd63t
//
type Bucket interface {
	// Retrieve data for the object with the given key.
	GetObject(key string) (data []byte, err error)

	// Store the supplied data with the given key, overwriting any previous
	// version. The object is created with the default ACL of "private".
	StoreObject(key string, data []byte) error
}

// OpenBucket returns a Bucket tied to a given name in whe given region. You
// must have previously created the bucket in the region, and the supplied
// access key must have access to it.
//
// To easily create a bucket, use the AWS Console:
//
//     http://aws.amazon.com/console/
//
func OpenBucket(name string, region Region, key aws.AccessKey) (Bucket, error) {
	// Create a connection to the given region's endpoint.
	endpoint := &url.URL{Scheme: "https", Host: string(region)}
	httpConn, err := http.NewConn(endpoint)
	if err != nil {
		return nil, fmt.Errorf("http.NewConn: %v", err)
	}

	// Create an appropriate request signer.
	signer, err := auth.NewSigner(&key)
	if err != nil {
		return nil, fmt.Errorf("auth.NewSigner: %v", err)
	}

	return openBucket(name, httpConn, signer, time.RealClock())
}

// A version of OpenBucket with the ability to inject dependencies, for
// testability.
func openBucket(
	name string,
	httpConn http.Conn,
	signer auth.Signer,
	clock time.Clock) (Bucket, error) {
	return &bucket{name, httpConn, signer, clock}, nil
}

type bucket struct {
	name     string
	httpConn http.Conn
	signer   auth.Signer
	clock    time.Clock
}

func (b *bucket) GetObject(key string) (data []byte, err error) {
	return nil, fmt.Errorf("TODO: Implement bucket.GetObject.")
}

func (b *bucket) StoreObject(key string, data []byte) error {
	// Validate the key.
	if len(key) > 1024 {
		return fmt.Errorf("Keys may be no longer than 1024 bytes.")
	}

	if !utf8.ValidString(key) {
		return fmt.Errorf("Keys must be valid UTF-8.")
	}

	// Build an appropriate HTTP request.
	//
	// Reference:
	//     http://docs.amazonwebservices.com/AmazonS3/latest/API/RESTObjectPUT.html
	httpReq := &http.Request{
		Verb: "PUT",
		Path: fmt.Sprintf("/%s/%s", b.name, key),
		Body: data,
		Headers: map[string]string{
			"Date": b.clock.Now().UTC().Format(sys_time.RFC1123),
		},
	}

	// Sign the request.
	if err := b.signer.Sign(httpReq); err != nil {
		return fmt.Errorf("Sign: %v", err)
	}

	// Send the request.
	httpResp, err := b.httpConn.SendRequest(httpReq)
	if err != nil {
		return fmt.Errorf("SendRequest: %v", err)
	}

	// Check the response.
	if httpResp.StatusCode != 200 {
		return fmt.Errorf("Error from server: %d %s", httpResp.StatusCode, httpResp.Body)
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"archive/tar"
	"bytes"
	"io"
	"log"
	"net"
	"os/exec"
	"path"
	"strings"
	"time"

	"github.com/rneugeba/virtsock/go/vsock"
	"github.com/rneugeba/virtsock/go/hvsock"
)

func run(timeout time.Duration, w *tar.Writer, command string, args ...string) {
	log.Printf("Running %s", command)
	c := exec.Command(command, args...)
	stdoutPipe, err := c.StdoutPipe()
	if err != nil {
		log.Fatalf("Failed to create stdout pipe: %#v", err)
	}
	stderrPipe, err := c.StderrPipe()
	if err != nil {
		log.Fatalf("Failed to create stderr pipe: %#v", err)
	}
	var stdoutBuffer bytes.Buffer
	var stderrBuffer bytes.Buffer
	done := make(chan int)
	go func() {
		io.Copy(&stdoutBuffer, stdoutPipe)
		done <- 0
	}()
	go func() {
		io.Copy(&stderrBuffer, stderrPipe)
		done <- 0
	}()
	var timer *time.Timer
	timer = time.AfterFunc(timeout, func() {
		timer.Stop()
		if c.Process != nil {
			c.Process.Kill()
		}
	})
	_ = c.Run()
	<-done
	<-done
	timer.Stop()

	name := strings.Join(append([]string{path.Base(command)}, args...), " ")

	hdr := &tar.Header{
		Name: name + ".stdout",
		Mode: 0644,
		Size: int64(stdoutBuffer.Len()),
	}
	if err = w.WriteHeader(hdr); err != nil {
		log.Fatalln(err)
	}
	if _, err = w.Write(stdoutBuffer.Bytes()); err != nil {
		log.Fatalln(err)
	}
	hdr = &tar.Header{
		Name: name + ".stderr",
		Mode: 0644,
		Size: int64(stderrBuffer.Len()),
	}
	if err = w.WriteHeader(hdr); err != nil {
		log.Fatalln(err)
	}
	if _, err = w.Write(stderrBuffer.Bytes()); err != nil {
		log.Fatalln(err)
	}

}

func capture(w *tar.Writer) {
	t := 2 * time.Second

	run(t, w, "/bin/date")
	run(t, w, "/bin/uname", "-a")
	run(t, w, "/bin/ps", "uax")
	run(t, w, "/bin/netstat", "-tulpn")
	run(t, w, "/sbin/iptables-save")
	run(t, w, "/sbin/ifconfig", "-a")
	run(t, w, "/sbin/route", "-n")
	run(t, w, "/usr/sbin/brctl", "show")
	run(t, w, "/bin/dmesg")
	run(t, w, "/usr/bin/docker", "ps")
	run(t, w, "/usr/bin/tail", "/var/log/docker.log")
	run(t, w, "/bin/mount")
	run(t, w, "/bin/df")
	run(t, w, "/bin/ls", "-l", "/var")
	run(t, w, "/bin/ls", "-l", "/var/lib")
	run(t, w, "/bin/ls", "-l", "/var/lib/docker")
	run(t, w, "/usr/bin/diagnostics")
	run(t, w, "/bin/ping", "-w", "5", "8.8.8.8")
	run(t, w, "/bin/cp", "/etc/resolv.conf", ".")
	run(t, w, "/usr/bin/dig", "docker.com")
	run(t, w, "/usr/bin/wget", "-O", "-", "http://www.docker.com/")
}

func main() {
	listeners := make([]net.Listener, 0)

	ip, err := net.Listen("tcp", ":62374")
	if err != nil {
		log.Printf("Failed to bind to TCP port 62374: %#v", err)
	} else {
		listeners = append(listeners, ip)
	}
	vsock, err := vsock.Listen(uint(62374))
	if err != nil {
		log.Printf("Failed to bind to vsock port 62374: %#v", err)
	} else {
		listeners = append(listeners, vsock)
	}
	svcid, _ := hvsock.GuidFromString("445BA2CB-E69B-4912-8B42-D7F494D007EA")
	hvsock, err := hvsock.Listen(hvsock.HypervAddr{VmId: hvsock.GUID_WILDCARD, ServiceId: svcid})
	if err != nil {
		log.Printf("Failed to bind to hvsock port: %#v", err)
	} else {
		listeners = append(listeners, hvsock)
	}

	for _, l := range listeners {
		go func(l net.Listener) {
			for {
				conn, err := l.Accept()
				if err != nil {
					log.Printf("Error accepting connection: %#v", err)
					return // no more listening
				}
				go func(conn net.Conn) {
					w := tar.NewWriter(conn)
					capture(w)
					if err := w.Close(); err != nil {
						log.Println(err)
					}
					conn.Close()
				}(conn)
			}
		}(l)
	}
	forever := make(chan int)
	<-forever
}
<commit_msg>diagnostics: add more log files and more content from logfiles<commit_after>package main

import (
	"archive/tar"
	"bytes"
	"io"
	"log"
	"net"
	"os/exec"
	"path"
	"strings"
	"time"

	"github.com/rneugeba/virtsock/go/vsock"
	"github.com/rneugeba/virtsock/go/hvsock"
)

func run(timeout time.Duration, w *tar.Writer, command string, args ...string) {
	log.Printf("Running %s", command)
	c := exec.Command(command, args...)
	stdoutPipe, err := c.StdoutPipe()
	if err != nil {
		log.Fatalf("Failed to create stdout pipe: %#v", err)
	}
	stderrPipe, err := c.StderrPipe()
	if err != nil {
		log.Fatalf("Failed to create stderr pipe: %#v", err)
	}
	var stdoutBuffer bytes.Buffer
	var stderrBuffer bytes.Buffer
	done := make(chan int)
	go func() {
		io.Copy(&stdoutBuffer, stdoutPipe)
		done <- 0
	}()
	go func() {
		io.Copy(&stderrBuffer, stderrPipe)
		done <- 0
	}()
	var timer *time.Timer
	timer = time.AfterFunc(timeout, func() {
		timer.Stop()
		if c.Process != nil {
			c.Process.Kill()
		}
	})
	_ = c.Run()
	<-done
	<-done
	timer.Stop()

	name := strings.Join(append([]string{path.Base(command)}, args...), " ")

	hdr := &tar.Header{
		Name: name + ".stdout",
		Mode: 0644,
		Size: int64(stdoutBuffer.Len()),
	}
	if err = w.WriteHeader(hdr); err != nil {
		log.Fatalln(err)
	}
	if _, err = w.Write(stdoutBuffer.Bytes()); err != nil {
		log.Fatalln(err)
	}
	hdr = &tar.Header{
		Name: name + ".stderr",
		Mode: 0644,
		Size: int64(stderrBuffer.Len()),
	}
	if err = w.WriteHeader(hdr); err != nil {
		log.Fatalln(err)
	}
	if _, err = w.Write(stderrBuffer.Bytes()); err != nil {
		log.Fatalln(err)
	}

}

func capture(w *tar.Writer) {
	t := 2 * time.Second

	run(t, w, "/bin/date")
	run(t, w, "/bin/uname", "-a")
	run(t, w, "/bin/ps", "uax")
	run(t, w, "/bin/netstat", "-tulpn")
	run(t, w, "/sbin/iptables-save")
	run(t, w, "/sbin/ifconfig", "-a")
	run(t, w, "/sbin/route", "-n")
	run(t, w, "/usr/sbin/brctl", "show")
	run(t, w, "/bin/dmesg")
	run(t, w, "/usr/bin/docker", "ps")
	run(t, w, "/usr/bin/tail", "-100", "/var/log/docker.log")
	run(t, w, "/usr/bin/tail", "-100", "/var/log/messages")
	run(t, w, "/usr/bin/tail", "-100", "/var/log/proxy-vsockd.log")
	run(t, w, "/usr/bin/tail", "-100", "/var/log/vsudd.log")
	run(t, w, "/bin/mount")
	run(t, w, "/bin/df")
	run(t, w, "/bin/ls", "-l", "/var")
	run(t, w, "/bin/ls", "-l", "/var/lib")
	run(t, w, "/bin/ls", "-l", "/var/lib/docker")
	run(t, w, "/usr/bin/diagnostics")
	run(t, w, "/bin/ping", "-w", "5", "8.8.8.8")
	run(t, w, "/bin/cp", "/etc/resolv.conf", ".")
	run(t, w, "/usr/bin/dig", "docker.com")
	run(t, w, "/usr/bin/wget", "-O", "-", "http://www.docker.com/")
}

func main() {
	listeners := make([]net.Listener, 0)

	ip, err := net.Listen("tcp", ":62374")
	if err != nil {
		log.Printf("Failed to bind to TCP port 62374: %#v", err)
	} else {
		listeners = append(listeners, ip)
	}
	vsock, err := vsock.Listen(uint(62374))
	if err != nil {
		log.Printf("Failed to bind to vsock port 62374: %#v", err)
	} else {
		listeners = append(listeners, vsock)
	}
	svcid, _ := hvsock.GuidFromString("445BA2CB-E69B-4912-8B42-D7F494D007EA")
	hvsock, err := hvsock.Listen(hvsock.HypervAddr{VmId: hvsock.GUID_WILDCARD, ServiceId: svcid})
	if err != nil {
		log.Printf("Failed to bind to hvsock port: %#v", err)
	} else {
		listeners = append(listeners, hvsock)
	}

	for _, l := range listeners {
		go func(l net.Listener) {
			for {
				conn, err := l.Accept()
				if err != nil {
					log.Printf("Error accepting connection: %#v", err)
					return // no more listening
				}
				go func(conn net.Conn) {
					w := tar.NewWriter(conn)
					capture(w)
					if err := w.Close(); err != nil {
						log.Println(err)
					}
					conn.Close()
				}(conn)
			}
		}(l)
	}
	forever := make(chan int)
	<-forever
}
<|endoftext|>


<commit_before>// The MIT License (MIT)
//
// Copyright (c) 2014, 2015 Jamie Alquiza
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
package main

import (
	"flag"
	"fmt"
	"log"
	"math"
	"math/rand"
	"os"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

	"github.com/jamiealquiza/tachymeter"
	"gopkg.in/Shopify/sarama.v1"
)

type config struct {
	brokers          []string
	topic            string
	msgSize          int
	msgRate          uint64
	batchSize        int
	compression      sarama.CompressionCodec
	workers          int
	writersPerWorker int
	noop             bool
}

var (
	Config = &config{}

	// Character selection for random messages.
	chars = []byte("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890!@#$^&*(){}][:<>.")

	// Counters / misc.
	sentCnt uint64
)

func init() {
	flag.StringVar(&Config.topic, "topic", "sangrenel", "Kafka topic to produce to")
	flag.IntVar(&Config.msgSize, "message-size", 300, "Message size (bytes)")
	flag.Uint64Var(&Config.msgRate, "produce-rate", 100000000, "Global write rate limit (messages/sec)")
	flag.IntVar(&Config.batchSize, "message-batch-size", 1, "Messages per batch")
	compression := flag.String("compression", "none", "Message compression: none, gzip, snappy")
	flag.BoolVar(&Config.noop, "noop", false, "Test message generation performance (does not connect to Kafka)")
	flag.IntVar(&Config.workers, "workers", 1, "Number of workers")
	flag.IntVar(&Config.writersPerWorker, "writers-per-worker", 5, "Number of writer (Kafka producer) goroutines per worker")
	brokerString := flag.String("brokers", "localhost:9092", "Comma delimited list of Kafka brokers")
	flag.Parse()

	Config.brokers = strings.Split(*brokerString, ",")

	switch *compression {
	case "gzip":
		Config.compression = sarama.CompressionGZIP
	case "snappy":
		Config.compression = sarama.CompressionSnappy
	case "none":
		Config.compression = sarama.CompressionNone
	default:
		fmt.Printf("Invalid compression option: %s\n", *compression)
		os.Exit(1)
	}
}

func main() {
	if graphiteIp != "" {
		go graphiteWriter()
	}

	// Print Sangrenel startup info.
	fmt.Printf("\nStarting %d client workers, %d writers per worker\n", Config.workers, Config.writersPerWorker)
	fmt.Printf("Message size %d bytes, %d message limit per batch\n", Config.msgSize, Config.batchSize)

	switch Config.compression {
	case sarama.CompressionNone:
		fmt.Println("Compression: none")
	case sarama.CompressionGZIP:
		fmt.Println("Compression: GZIP")
	case sarama.CompressionSnappy:
		fmt.Println("Compression: Snappy")
	}

	t := tachymeter.New(&tachymeter.Config{Size: 300000, Safe: true})

	// Start client workers.
	for i := 0; i < Config.workers; i++ {
		go worker(i+1, t)
	}

	var currCnt, lastCnt uint64
	interval := 5 * time.Second
	ticker := time.Tick(interval)
	start := time.Now()

	for {
		<-ticker

		intervalTime := time.Since(start).Seconds()

		// Set tachymeter wall time.
		t.SetWallTime(time.Since(start))

		// Get the sent count from the last interval, then the delta
		// (sentSinceLastInterval) between the current and last interval.
		lastCnt = currCnt
		currCnt = atomic.LoadUint64(&sentCnt)
		sentSinceLastInterval := currCnt - lastCnt

		outputBytes, outputString := calcOutput(intervalTime, sentSinceLastInterval)

		// Summarize tachymeter data.
		stats := t.Calc()

		// Update the metrics map for the Graphite writer.
		metrics["rate"] = float64(sentSinceLastInterval) / intervalTime
		metrics["output"] = outputBytes
		metrics["p99"] = (float64(stats.Time.P99.Nanoseconds()) / 1000) / 1000
		metrics["timestamp"] = float64(time.Now().Unix())
		// Ship metrics if configured.
		if graphiteIp != "" {
			metricsOutgoing <- metrics
		}

		// Write output stats.
		fmt.Println()
		log.Printf("Generating %s @ %.0f messages/sec | topic: %s | %.2fms p99 latency\n",
			outputString,
			metrics["rate"],
			Config.topic,
			metrics["p99"])

		fmt.Printf("[Batch Statistics, last %.1fs]\n", intervalTime)
		stats.Dump()

		// Check if the tacymeter size needs to be increased
		// to avoid sampling. Otherwise, just reset it.
		if int(sentSinceLastInterval) > len(t.Times) {
			newTachy := tachymeter.New(&tachymeter.Config{Size: int(2 * sentSinceLastInterval), Safe: true})
			// This is actually dangerous;
			// this could swap in a tachy with unlocked
			// mutexes while the current one has locks held.
			*t = *newTachy
		} else {
			t.Reset()
		}

		// Reset interval time.
		start = time.Now()
	}
}

// worker is a high level producer unit and holds a single
// Kafka client. The worker's Kafka client is shared by n (Config.writersPerWorker)
// writer instances that perform the message generation and writing.
func worker(n int, t *tachymeter.Tachymeter) {
	switch Config.noop {
	case false:
		cId := "worker_" + strconv.Itoa(n)

		conf := sarama.NewConfig()
		conf.Producer.Compression = Config.compression
		conf.Producer.Return.Successes = true
		conf.Producer.Flush.MaxMessages = Config.batchSize
		conf.Producer.MaxMessageBytes = Config.msgSize + 50

		client, err := sarama.NewClient(Config.brokers, conf)
		if err != nil {
			log.Println(err)
			os.Exit(1)
		} else {
			log.Printf("%s connected\n", cId)
		}

		for i := 0; i < Config.writersPerWorker; i++ {
			go writer(client, t)
		}
	// If noop, we're not creating connections at all.
	// Just generate messages and burn CPU.
	default:
		for i := 0; i < Config.writersPerWorker; i++ {
			go dummyWriter(t)
		}
	}

	wait := make(chan bool)
	<-wait
}

// writer generates random messages and write to Kafka.
// Each wrtier belongs to a parent worker. Writers
// throttle writes according to a global rate limiter
// and report write throughput statistics up through
// a shared tachymeter.
func writer(c sarama.Client, t *tachymeter.Tachymeter) {
	// Init the producer.
	producer, err := sarama.NewSyncProducerFromClient(c)
	if err != nil {
		log.Println(err.Error())
	}
	defer producer.Close()

	source := rand.NewSource(time.Now().UnixNano())
	generator := rand.New(source)
	msgBatch := make([]*sarama.ProducerMessage, 0, Config.batchSize)

	for {
		// Message rate limiting works by having all writer loops incrementing
		// a global counter and tracking the aggregate per-second progress.
		// If the configured rate is met, the worker will sleep
		// for the remainder of the 1 second window.
		rateEnd := time.Now().Add(time.Second)
		countStart := atomic.LoadUint64(&sentCnt)
		var start time.Time // TODO revisit if this should be moved.

		for {
			// Break if the global rate limit was met, or, if
			// we'd exceed it assuming all writers wrote a max batch size
			// for this interval.
			intervalSent := atomic.LoadUint64(&sentCnt) - countStart
			sendEstimate := intervalSent + uint64(Config.batchSize*Config.workers*(Config.writersPerWorker-1))
			if sendEstimate >= Config.msgRate {
				break
			}

			// Estimate the batch size. This should shrink
			// if we're near the rate limit. Estimated batch size =
			// amount left to send for this interval / number of writers
			// we have available to send this amount. If the estimate
			// is lower than the configured batch size, send that amount
			// instead.
			toSend := (Config.msgRate - intervalSent) / uint64((Config.workers * Config.writersPerWorker))
			n := int(math.Min(float64(toSend), float64(Config.batchSize)))

			for i := 0; i < n; i++ {
				// Gen message.
				msgData := make([]byte, Config.msgSize)
				randMsg(msgData, *generator)
				msg := &sarama.ProducerMessage{Topic: Config.topic, Value: sarama.ByteEncoder(msgData)}
				// Append to batch.
				msgBatch = append(msgBatch, msg)
			}

			start = time.Now()
			err = producer.SendMessages(msgBatch)
			if err != nil {
				log.Println(err)
			} else {
				t.AddTime(time.Since(start))
				atomic.AddUint64(&sentCnt, uint64(len(msgBatch)))
			}

			msgBatch = msgBatch[:0]
		}

		// If the global per-second rate limit was met,
		// the inner loop breaks and the outer loop sleeps for the interval remainder.
		time.Sleep(rateEnd.Sub(time.Now()) + time.Since(start))
	}
}

// dummyWriter is initialized by the worker(s) if Config.noop is True.
// dummyWriter performs the message generation step of the normal writer,
// but doesn't connect to / attempt to send anything to Kafka. This is used
// purely for testing message generation performance.
func dummyWriter(t *tachymeter.Tachymeter) {
	source := rand.NewSource(time.Now().UnixNano())
	generator := rand.New(source)
	msg := make([]byte, Config.msgSize)

	var sent int64

	for {
		randMsg(msg, *generator)

		t.AddTime(time.Duration(0))
		sent++
		if sent == 10 {
			atomic.AddUint64(&sentCnt, 10)
			sent = 0
		}
	}
}

// randMsg returns a random message generated from the chars byte slice.
// Message length of m bytes as defined by Config.msgSize.
func randMsg(m []byte, generator rand.Rand) {
	for i := range m {
		m[i] = chars[generator.Intn(len(chars))]
	}
}

// calcOutput takes a duration t and messages sent
// and returns message rates in human readable network speeds.
func calcOutput(t float64, n uint64) (float64, string) {
	m := (float64(n) / t) * float64(Config.msgSize)
	var o string
	switch {
	case m >= 131072:
		o = strconv.FormatFloat(m/131072, 'f', 0, 64) + "Mb/sec"
	case m < 131072:
		o = strconv.FormatFloat(m/1024, 'f', 0, 64) + "KB/sec"
	}
	return m, o
}
<commit_msg>updated timers, output<commit_after>// The MIT License (MIT)
//
// Copyright (c) 2014, 2015 Jamie Alquiza
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
package main

import (
	"flag"
	"fmt"
	"log"
	"math"
	"math/rand"
	"os"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

	"github.com/jamiealquiza/tachymeter"
	"gopkg.in/Shopify/sarama.v1"
)

type config struct {
	brokers          []string
	topic            string
	msgSize          int
	msgRate          uint64
	batchSize        int
	compression      sarama.CompressionCodec
	workers          int
	writersPerWorker int
	noop             bool
}

var (
	Config = &config{}

	// Character selection for random messages.
	chars = []byte("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890!@#$^&*(){}][:<>.")

	// Counters / misc.
	sentCnt uint64
)

func init() {
	flag.StringVar(&Config.topic, "topic", "sangrenel", "Kafka topic to produce to")
	flag.IntVar(&Config.msgSize, "message-size", 300, "Message size (bytes)")
	flag.Uint64Var(&Config.msgRate, "produce-rate", 100000000, "Global write rate limit (messages/sec)")
	flag.IntVar(&Config.batchSize, "message-batch-size", 1, "Messages per batch")
	compression := flag.String("compression", "none", "Message compression: none, gzip, snappy")
	flag.BoolVar(&Config.noop, "noop", false, "Test message generation performance (does not connect to Kafka)")
	flag.IntVar(&Config.workers, "workers", 1, "Number of workers")
	flag.IntVar(&Config.writersPerWorker, "writers-per-worker", 5, "Number of writer (Kafka producer) goroutines per worker")
	brokerString := flag.String("brokers", "localhost:9092", "Comma delimited list of Kafka brokers")
	flag.Parse()

	Config.brokers = strings.Split(*brokerString, ",")

	switch *compression {
	case "gzip":
		Config.compression = sarama.CompressionGZIP
	case "snappy":
		Config.compression = sarama.CompressionSnappy
	case "none":
		Config.compression = sarama.CompressionNone
	default:
		fmt.Printf("Invalid compression option: %s\n", *compression)
		os.Exit(1)
	}
}

func main() {
	if graphiteIp != "" {
		go graphiteWriter()
	}

	// Print Sangrenel startup info.
	fmt.Printf("\nStarting %d client workers, %d writers per worker\n", Config.workers, Config.writersPerWorker)
	fmt.Printf("Message size %d bytes, %d message limit per batch\n", Config.msgSize, Config.batchSize)

	switch Config.compression {
	case sarama.CompressionNone:
		fmt.Println("Compression: none")
	case sarama.CompressionGZIP:
		fmt.Println("Compression: GZIP")
	case sarama.CompressionSnappy:
		fmt.Println("Compression: Snappy")
	}

	t := tachymeter.New(&tachymeter.Config{Size: 300000, Safe: true})

	// Start client workers.
	for i := 0; i < Config.workers; i++ {
		go worker(i+1, t)
	}

	var currCnt, lastCnt uint64
	interval := 5 * time.Second
	ticker := time.Tick(interval)
	start := time.Now()

	for {
		<-ticker

		intervalTime := time.Since(start).Seconds()

		// Set tachymeter wall time.
		t.SetWallTime(time.Since(start))

		// Get the sent count from the last interval, then the delta
		// (sentSinceLastInterval) between the current and last interval.
		lastCnt = currCnt
		currCnt = atomic.LoadUint64(&sentCnt)
		sentSinceLastInterval := currCnt - lastCnt

		outputBytes, outputString := calcOutput(intervalTime, sentSinceLastInterval)

		// Summarize tachymeter data.
		stats := t.Calc()

		// Update the metrics map for the Graphite writer.
		metrics["rate"] = float64(sentSinceLastInterval) / intervalTime
		metrics["output"] = outputBytes
		metrics["p99"] = (float64(stats.Time.P99.Nanoseconds()) / 1000) / 1000
		metrics["timestamp"] = float64(time.Now().Unix())
		// Ship metrics if configured.
		if graphiteIp != "" {
			metricsOutgoing <- metrics
		}

		// Write output stats.
		fmt.Println()
		log.Printf("Generating %s @ %.0f messages/sec | topic: %s | %.2fms p99 batch latency\n",
			outputString,
			metrics["rate"],
			Config.topic,
			metrics["p99"])

		fmt.Printf("> Batch Statistics, Last %.1fs:\n", intervalTime)
		stats.Dump()

		// Check if the tacymeter size needs to be increased
		// to avoid sampling. Otherwise, just reset it.
		if int(sentSinceLastInterval) > len(t.Times) {
			newTachy := tachymeter.New(&tachymeter.Config{Size: int(2 * sentSinceLastInterval), Safe: true})
			// This is actually dangerous;
			// this could swap in a tachy with unlocked
			// mutexes while the current one has locks held.
			*t = *newTachy
		} else {
			t.Reset()
		}

		// Reset interval time.
		start = time.Now()
	}
}

// worker is a high level producer unit and holds a single
// Kafka client. The worker's Kafka client is shared by n (Config.writersPerWorker)
// writer instances that perform the message generation and writing.
func worker(n int, t *tachymeter.Tachymeter) {
	switch Config.noop {
	case false:
		cId := "worker_" + strconv.Itoa(n)

		conf := sarama.NewConfig()
		conf.Producer.Compression = Config.compression
		conf.Producer.Return.Successes = true
		conf.Producer.Flush.MaxMessages = Config.batchSize
		conf.Producer.MaxMessageBytes = Config.msgSize + 50

		client, err := sarama.NewClient(Config.brokers, conf)
		if err != nil {
			log.Println(err)
			os.Exit(1)
		} else {
			log.Printf("%s connected\n", cId)
		}

		for i := 0; i < Config.writersPerWorker; i++ {
			go writer(client, t)
		}
	// If noop, we're not creating connections at all.
	// Just generate messages and burn CPU.
	default:
		for i := 0; i < Config.writersPerWorker; i++ {
			go dummyWriter(t)
		}
	}

	wait := make(chan bool)
	<-wait
}

// writer generates random messages and write to Kafka.
// Each wrtier belongs to a parent worker. Writers
// throttle writes according to a global rate limiter
// and report write throughput statistics up through
// a shared tachymeter.
func writer(c sarama.Client, t *tachymeter.Tachymeter) {
	// Init the producer.
	producer, err := sarama.NewSyncProducerFromClient(c)
	if err != nil {
		log.Println(err.Error())
	}
	defer producer.Close()

	source := rand.NewSource(time.Now().UnixNano())
	generator := rand.New(source)
	msgBatch := make([]*sarama.ProducerMessage, 0, Config.batchSize)

	for {
		// Message rate limiting works by having all writer loops incrementing
		// a global counter and tracking the aggregate per-second progress.
		// If the configured rate is met, the worker will sleep
		// for the remainder of the 1 second window.
		intervalEnd := time.Now().Add(time.Second)
		countStart := atomic.LoadUint64(&sentCnt)

		var sendTime time.Time

		for {
			// Break if the global rate limit was met, or, if
			// we'd exceed it assuming all writers wrote a max batch size
			// for this interval.
			intervalSent := atomic.LoadUint64(&sentCnt) - countStart
			sendEstimate := intervalSent + uint64(Config.batchSize*Config.workers*(Config.writersPerWorker-1))
			if sendEstimate >= Config.msgRate {
				break
			}

			// Estimate the batch size. This should shrink
			// if we're near the rate limit. Estimated batch size =
			// amount left to send for this interval / number of writers
			// we have available to send this amount. If the estimate
			// is lower than the configured batch size, send that amount
			// instead.
			toSend := (Config.msgRate - intervalSent) / uint64((Config.workers * Config.writersPerWorker))
			n := int(math.Min(float64(toSend), float64(Config.batchSize)))

			for i := 0; i < n; i++ {
				// Gen message.
				msgData := make([]byte, Config.msgSize)
				randMsg(msgData, *generator)
				msg := &sarama.ProducerMessage{Topic: Config.topic, Value: sarama.ByteEncoder(msgData)}
				// Append to batch.
				msgBatch = append(msgBatch, msg)
			}

			sendTime = time.Now()
			err = producer.SendMessages(msgBatch)
			if err != nil {
				log.Println(err)
			} else {
				t.AddTime(time.Since(sendTime))
				atomic.AddUint64(&sentCnt, uint64(len(msgBatch)))
			}

			msgBatch = msgBatch[:0]
		}

		// If the global per-second rate limit was met,
		// the inner loop breaks and the outer loop sleeps for the interval remainder.
		time.Sleep(intervalEnd.Sub(time.Now()))
	}
}

// dummyWriter is initialized by the worker(s) if Config.noop is True.
// dummyWriter performs the message generation step of the normal writer,
// but doesn't connect to / attempt to send anything to Kafka. This is used
// purely for testing message generation performance.
func dummyWriter(t *tachymeter.Tachymeter) {
	source := rand.NewSource(time.Now().UnixNano())
	generator := rand.New(source)
	msg := make([]byte, Config.msgSize)

	var sent int64

	for {
		randMsg(msg, *generator)

		t.AddTime(time.Duration(0))
		sent++
		if sent == 10 {
			atomic.AddUint64(&sentCnt, 10)
			sent = 0
		}
	}
}

// randMsg returns a random message generated from the chars byte slice.
// Message length of m bytes as defined by Config.msgSize.
func randMsg(m []byte, generator rand.Rand) {
	for i := range m {
		m[i] = chars[generator.Intn(len(chars))]
	}
}

// calcOutput takes a duration t and messages sent
// and returns message rates in human readable network speeds.
func calcOutput(t float64, n uint64) (float64, string) {
	m := (float64(n) / t) * float64(Config.msgSize)
	var o string
	switch {
	case m >= 131072:
		o = strconv.FormatFloat(m/131072, 'f', 0, 64) + "Mb/sec"
	case m < 131072:
		o = strconv.FormatFloat(m/1024, 'f', 0, 64) + "KB/sec"
	}
	return m, o
}
<|endoftext|>


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

package scorch

import (
	"fmt"
	"io/ioutil"
	"os"
	"sync"

	"github.com/RoaringBitmap/roaring"
	"github.com/blevesearch/bleve/document"
	"github.com/blevesearch/bleve/index"
	"github.com/blevesearch/bleve/index/scorch/segment"
	bolt "go.etcd.io/bbolt"
)

const DefaultBuilderBatchSize = 1000
const DefaultBuilderMergeMax = 10

type Builder struct {
	m         sync.Mutex
	segCount  uint64
	path      string
	buildPath string
	segPaths  []string
	batchSize int
	mergeMax  int
	batch     *index.Batch
	internal  map[string][]byte
	segPlugin segment.Plugin
}

func NewBuilder(config map[string]interface{}) (*Builder, error) {
	path, ok := config["path"].(string)
	if !ok {
		return nil, fmt.Errorf("must specify path")
	}

	buildPathPrefix, _ := config["buildPathPrefix"].(string)
	buildPath, err := ioutil.TempDir(buildPathPrefix, "scorch-offline-build")
	if err != nil {
		return nil, err
	}

	rv := &Builder{
		path:      path,
		buildPath: buildPath,
		mergeMax:  DefaultBuilderMergeMax,
		batchSize: DefaultBuilderBatchSize,
		batch:     index.NewBatch(),
		segPlugin: defaultSegmentPlugin,
	}

	err = rv.parseConfig(config)
	if err != nil {
		return nil, fmt.Errorf("error parsing builder config: %v", err)
	}

	return rv, nil
}

func (o *Builder) parseConfig(config map[string]interface{}) (err error) {
	if v, ok := config["mergeMax"]; ok {
		var t int
		if t, err = parseToInteger(v); err != nil {
			return fmt.Errorf("mergeMax parse err: %v", err)
		}
		if t > 0 {
			o.mergeMax = t
		}
	}

	if v, ok := config["batchSize"]; ok {
		var t int
		if t, err = parseToInteger(v); err != nil {
			return fmt.Errorf("batchSize parse err: %v", err)
		}
		if t > 0 {
			o.batchSize = t
		}
	}

	if v, ok := config["internal"]; ok {
		if vinternal, ok := v.(map[string][]byte); ok {
			o.internal = vinternal
		}
	}

	forcedSegmentType, forcedSegmentVersion, err := configForceSegmentTypeVersion(config)
	if err != nil {
		return err
	}
	if forcedSegmentType != "" && forcedSegmentVersion != 0 {
		segPlugin, err := chooseSegmentPlugin(forcedSegmentType,
			uint32(forcedSegmentVersion))
		if err != nil {
			o.segPlugin = segPlugin
		}
	}

	return nil
}

// Index will place the document into the index.
// It is invalid to index the same document multiple times.
func (o *Builder) Index(doc *document.Document) error {
	o.m.Lock()
	defer o.m.Unlock()

	o.batch.Update(doc)

	return o.maybeFlushBatchLOCKED(o.batchSize)
}

func (o *Builder) maybeFlushBatchLOCKED(moreThan int) error {
	if len(o.batch.IndexOps) >= moreThan {
		defer o.batch.Reset()
		return o.executeBatchLOCKED(o.batch)
	}
	return nil
}

func (o *Builder) executeBatchLOCKED(batch *index.Batch) (err error) {
	analysisResults := make([]*index.AnalysisResult, 0, len(batch.IndexOps))
	for _, doc := range batch.IndexOps {
		if doc != nil {
			// insert _id field
			doc.AddField(document.NewTextFieldCustom("_id", nil, []byte(doc.ID), document.IndexField|document.StoreField, nil))
			// perform analysis directly
			analysisResult := analyze(doc)
			analysisResults = append(analysisResults, analysisResult)
		}
	}

	seg, _, err := o.segPlugin.New(analysisResults)
	if err != nil {
		return fmt.Errorf("error building segment base: %v", err)
	}

	filename := zapFileName(o.segCount)
	o.segCount++
	path := o.buildPath + string(os.PathSeparator) + filename

	if segUnpersisted, ok := seg.(segment.UnpersistedSegment); ok {
		err = segUnpersisted.Persist(path)
		if err != nil {
			return fmt.Errorf("error persisting segment base to %s: %v", path, err)
		}

		o.segPaths = append(o.segPaths, path)
		return nil
	}

	return fmt.Errorf("new segment does not implement unpersisted: %T", seg)
}

func (o *Builder) doMerge() error {
	// as long as we have more than 1 segment, keep merging
	for len(o.segPaths) > 1 {

		// merge the next <mergeMax> number of segments into one new one
		// or, if there are fewer than <mergeMax> remaining, merge them all
		mergeCount := o.mergeMax
		if mergeCount > len(o.segPaths) {
			mergeCount = len(o.segPaths)
		}

		mergePaths := o.segPaths[0:mergeCount]
		o.segPaths = o.segPaths[mergeCount:]

		// open each of the segments to be merged
		mergeSegs := make([]segment.Segment, 0, mergeCount)

		// closeOpenedSegs attempts to close all opened
		// segments even if an error occurs, in which case
		// the first error is returned
		closeOpenedSegs := func() error {
			var err error
			for _, seg := range mergeSegs {
				clErr := seg.Close()
				if clErr != nil && err == nil {
					err = clErr
				}
			}
			return err
		}

		for _, mergePath := range mergePaths {
			seg, err := o.segPlugin.Open(mergePath)
			if err != nil {
				_ = closeOpenedSegs()
				return fmt.Errorf("error opening segment (%s) for merge: %v", mergePath, err)
			}
			mergeSegs = append(mergeSegs, seg)
		}

		// do the merge
		mergedSegPath := o.buildPath + string(os.PathSeparator) + zapFileName(o.segCount)
		drops := make([]*roaring.Bitmap, mergeCount)
		_, _, err := o.segPlugin.Merge(mergeSegs, drops, mergedSegPath, nil, nil)
		if err != nil {
			_ = closeOpenedSegs()
			return fmt.Errorf("error merging segments (%v): %v", mergePaths, err)
		}
		o.segCount++
		o.segPaths = append(o.segPaths, mergedSegPath)

		// close segments opened for merge
		err = closeOpenedSegs()
		if err != nil {
			return fmt.Errorf("error closing opened segments: %v", err)
		}

		// remove merged segments
		for _, mergePath := range mergePaths {
			err = os.RemoveAll(mergePath)
			if err != nil {
				return fmt.Errorf("error removing segment %s after merge: %v", mergePath, err)
			}
		}
	}

	return nil
}

func (o *Builder) Close() error {
	o.m.Lock()
	defer o.m.Unlock()

	// see if there is a partial batch
	err := o.maybeFlushBatchLOCKED(1)
	if err != nil {
		return fmt.Errorf("error flushing batch before close: %v", err)
	}

	// perform all the merging
	err = o.doMerge()
	if err != nil {
		return fmt.Errorf("error while merging: %v", err)
	}

	// ensure the store path exists
	err = os.MkdirAll(o.path, 0700)
	if err != nil {
		return err
	}

	// move final segment into place
	// segment id 2 is chosen to match the behavior of a scorch
	// index which indexes a single batch of data
	finalSegPath := o.path + string(os.PathSeparator) + zapFileName(2)
	err = os.Rename(o.segPaths[0], finalSegPath)
	if err != nil {
		return fmt.Errorf("error moving final segment into place: %v", err)
	}

	// remove the buildPath, as it is no longer needed
	err = os.RemoveAll(o.buildPath)
	if err != nil {
		return fmt.Errorf("error removing build path: %v", err)
	}

	// prepare wrapping
	seg, err := o.segPlugin.Open(finalSegPath)
	if err != nil {
		return fmt.Errorf("error opening final segment")
	}

	// create a segment snapshot for this segment
	ss := &SegmentSnapshot{
		segment: seg,
	}
	is := &IndexSnapshot{
		epoch:    3, // chosen to match scorch behavior when indexing a single batch
		segment:  []*SegmentSnapshot{ss},
		creator:  "scorch-builder",
		internal: o.internal,
	}

	// create the root bolt
	rootBoltPath := o.path + string(os.PathSeparator) + "root.bolt"
	rootBolt, err := bolt.Open(rootBoltPath, 0600, nil)
	if err != nil {
		return err
	}

	// start a write transaction
	tx, err := rootBolt.Begin(true)
	if err != nil {
		return err
	}

	// fill the root bolt with this fake index snapshot
	_, _, err = prepareBoltSnapshot(is, tx, o.path, o.segPlugin)
	if err != nil {
		_ = tx.Rollback()
		_ = rootBolt.Close()
		return fmt.Errorf("error preparing bolt snapshot in root.bolt: %v", err)
	}

	// commit bolt data
	err = tx.Commit()
	if err != nil {
		_ = rootBolt.Close()
		return fmt.Errorf("error committing bolt tx in root.bolt: %v", err)
	}

	// close bolt
	err = rootBolt.Close()
	if err != nil {
		return fmt.Errorf("error closing root.bolt: %v", err)
	}

	// close final segment
	err = seg.Close()
	if err != nil {
		return fmt.Errorf("error closing final segment: %v", err)
	}
	return nil
}
<commit_msg>allow segment type/version override to work (#1407)<commit_after>//  Copyright (c) 2019 Couchbase, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// 		http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package scorch

import (
	"fmt"
	"io/ioutil"
	"os"
	"sync"

	"github.com/RoaringBitmap/roaring"
	"github.com/blevesearch/bleve/document"
	"github.com/blevesearch/bleve/index"
	"github.com/blevesearch/bleve/index/scorch/segment"
	bolt "go.etcd.io/bbolt"
)

const DefaultBuilderBatchSize = 1000
const DefaultBuilderMergeMax = 10

type Builder struct {
	m         sync.Mutex
	segCount  uint64
	path      string
	buildPath string
	segPaths  []string
	batchSize int
	mergeMax  int
	batch     *index.Batch
	internal  map[string][]byte
	segPlugin segment.Plugin
}

func NewBuilder(config map[string]interface{}) (*Builder, error) {
	path, ok := config["path"].(string)
	if !ok {
		return nil, fmt.Errorf("must specify path")
	}

	buildPathPrefix, _ := config["buildPathPrefix"].(string)
	buildPath, err := ioutil.TempDir(buildPathPrefix, "scorch-offline-build")
	if err != nil {
		return nil, err
	}

	rv := &Builder{
		path:      path,
		buildPath: buildPath,
		mergeMax:  DefaultBuilderMergeMax,
		batchSize: DefaultBuilderBatchSize,
		batch:     index.NewBatch(),
		segPlugin: defaultSegmentPlugin,
	}

	err = rv.parseConfig(config)
	if err != nil {
		return nil, fmt.Errorf("error parsing builder config: %v", err)
	}

	return rv, nil
}

func (o *Builder) parseConfig(config map[string]interface{}) (err error) {
	if v, ok := config["mergeMax"]; ok {
		var t int
		if t, err = parseToInteger(v); err != nil {
			return fmt.Errorf("mergeMax parse err: %v", err)
		}
		if t > 0 {
			o.mergeMax = t
		}
	}

	if v, ok := config["batchSize"]; ok {
		var t int
		if t, err = parseToInteger(v); err != nil {
			return fmt.Errorf("batchSize parse err: %v", err)
		}
		if t > 0 {
			o.batchSize = t
		}
	}

	if v, ok := config["internal"]; ok {
		if vinternal, ok := v.(map[string][]byte); ok {
			o.internal = vinternal
		}
	}

	forcedSegmentType, forcedSegmentVersion, err := configForceSegmentTypeVersion(config)
	if err != nil {
		return err
	}
	if forcedSegmentType != "" && forcedSegmentVersion != 0 {
		segPlugin, err := chooseSegmentPlugin(forcedSegmentType,
			uint32(forcedSegmentVersion))
		if err != nil {
			return err
		}
		o.segPlugin = segPlugin
	}

	return nil
}

// Index will place the document into the index.
// It is invalid to index the same document multiple times.
func (o *Builder) Index(doc *document.Document) error {
	o.m.Lock()
	defer o.m.Unlock()

	o.batch.Update(doc)

	return o.maybeFlushBatchLOCKED(o.batchSize)
}

func (o *Builder) maybeFlushBatchLOCKED(moreThan int) error {
	if len(o.batch.IndexOps) >= moreThan {
		defer o.batch.Reset()
		return o.executeBatchLOCKED(o.batch)
	}
	return nil
}

func (o *Builder) executeBatchLOCKED(batch *index.Batch) (err error) {
	analysisResults := make([]*index.AnalysisResult, 0, len(batch.IndexOps))
	for _, doc := range batch.IndexOps {
		if doc != nil {
			// insert _id field
			doc.AddField(document.NewTextFieldCustom("_id", nil, []byte(doc.ID), document.IndexField|document.StoreField, nil))
			// perform analysis directly
			analysisResult := analyze(doc)
			analysisResults = append(analysisResults, analysisResult)
		}
	}

	seg, _, err := o.segPlugin.New(analysisResults)
	if err != nil {
		return fmt.Errorf("error building segment base: %v", err)
	}

	filename := zapFileName(o.segCount)
	o.segCount++
	path := o.buildPath + string(os.PathSeparator) + filename

	if segUnpersisted, ok := seg.(segment.UnpersistedSegment); ok {
		err = segUnpersisted.Persist(path)
		if err != nil {
			return fmt.Errorf("error persisting segment base to %s: %v", path, err)
		}

		o.segPaths = append(o.segPaths, path)
		return nil
	}

	return fmt.Errorf("new segment does not implement unpersisted: %T", seg)
}

func (o *Builder) doMerge() error {
	// as long as we have more than 1 segment, keep merging
	for len(o.segPaths) > 1 {

		// merge the next <mergeMax> number of segments into one new one
		// or, if there are fewer than <mergeMax> remaining, merge them all
		mergeCount := o.mergeMax
		if mergeCount > len(o.segPaths) {
			mergeCount = len(o.segPaths)
		}

		mergePaths := o.segPaths[0:mergeCount]
		o.segPaths = o.segPaths[mergeCount:]

		// open each of the segments to be merged
		mergeSegs := make([]segment.Segment, 0, mergeCount)

		// closeOpenedSegs attempts to close all opened
		// segments even if an error occurs, in which case
		// the first error is returned
		closeOpenedSegs := func() error {
			var err error
			for _, seg := range mergeSegs {
				clErr := seg.Close()
				if clErr != nil && err == nil {
					err = clErr
				}
			}
			return err
		}

		for _, mergePath := range mergePaths {
			seg, err := o.segPlugin.Open(mergePath)
			if err != nil {
				_ = closeOpenedSegs()
				return fmt.Errorf("error opening segment (%s) for merge: %v", mergePath, err)
			}
			mergeSegs = append(mergeSegs, seg)
		}

		// do the merge
		mergedSegPath := o.buildPath + string(os.PathSeparator) + zapFileName(o.segCount)
		drops := make([]*roaring.Bitmap, mergeCount)
		_, _, err := o.segPlugin.Merge(mergeSegs, drops, mergedSegPath, nil, nil)
		if err != nil {
			_ = closeOpenedSegs()
			return fmt.Errorf("error merging segments (%v): %v", mergePaths, err)
		}
		o.segCount++
		o.segPaths = append(o.segPaths, mergedSegPath)

		// close segments opened for merge
		err = closeOpenedSegs()
		if err != nil {
			return fmt.Errorf("error closing opened segments: %v", err)
		}

		// remove merged segments
		for _, mergePath := range mergePaths {
			err = os.RemoveAll(mergePath)
			if err != nil {
				return fmt.Errorf("error removing segment %s after merge: %v", mergePath, err)
			}
		}
	}

	return nil
}

func (o *Builder) Close() error {
	o.m.Lock()
	defer o.m.Unlock()

	// see if there is a partial batch
	err := o.maybeFlushBatchLOCKED(1)
	if err != nil {
		return fmt.Errorf("error flushing batch before close: %v", err)
	}

	// perform all the merging
	err = o.doMerge()
	if err != nil {
		return fmt.Errorf("error while merging: %v", err)
	}

	// ensure the store path exists
	err = os.MkdirAll(o.path, 0700)
	if err != nil {
		return err
	}

	// move final segment into place
	// segment id 2 is chosen to match the behavior of a scorch
	// index which indexes a single batch of data
	finalSegPath := o.path + string(os.PathSeparator) + zapFileName(2)
	err = os.Rename(o.segPaths[0], finalSegPath)
	if err != nil {
		return fmt.Errorf("error moving final segment into place: %v", err)
	}

	// remove the buildPath, as it is no longer needed
	err = os.RemoveAll(o.buildPath)
	if err != nil {
		return fmt.Errorf("error removing build path: %v", err)
	}

	// prepare wrapping
	seg, err := o.segPlugin.Open(finalSegPath)
	if err != nil {
		return fmt.Errorf("error opening final segment")
	}

	// create a segment snapshot for this segment
	ss := &SegmentSnapshot{
		segment: seg,
	}
	is := &IndexSnapshot{
		epoch:    3, // chosen to match scorch behavior when indexing a single batch
		segment:  []*SegmentSnapshot{ss},
		creator:  "scorch-builder",
		internal: o.internal,
	}

	// create the root bolt
	rootBoltPath := o.path + string(os.PathSeparator) + "root.bolt"
	rootBolt, err := bolt.Open(rootBoltPath, 0600, nil)
	if err != nil {
		return err
	}

	// start a write transaction
	tx, err := rootBolt.Begin(true)
	if err != nil {
		return err
	}

	// fill the root bolt with this fake index snapshot
	_, _, err = prepareBoltSnapshot(is, tx, o.path, o.segPlugin)
	if err != nil {
		_ = tx.Rollback()
		_ = rootBolt.Close()
		return fmt.Errorf("error preparing bolt snapshot in root.bolt: %v", err)
	}

	// commit bolt data
	err = tx.Commit()
	if err != nil {
		_ = rootBolt.Close()
		return fmt.Errorf("error committing bolt tx in root.bolt: %v", err)
	}

	// close bolt
	err = rootBolt.Close()
	if err != nil {
		return fmt.Errorf("error closing root.bolt: %v", err)
	}

	// close final segment
	err = seg.Close()
	if err != nil {
		return fmt.Errorf("error closing final segment: %v", err)
	}
	return nil
}
<|endoftext|>


<commit_before>package config

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"strconv"
	"strings"

	"golang.org/x/crypto/acme/autocert"
)

// SecretStore represents a contract for a store capable of resolving secrets.
type SecretStore interface {
	Configure(config Config) error
	GetSecret(secretName string) (string, bool)
}

// CertificateStore represents a secret store which can be used for certificates.
type CertificateStore interface {
	autocert.Cache
}

// Provider represents a configurable provider.
type Provider interface {
	Name() string
	Configure(config map[string]interface{}) error
}

// Config represents a configuration interface.
type Config interface {
	Vault() *VaultConfig
}

// TLSConfig represents TLS listener configuration.
type TLSConfig struct {
	ListenAddr string `json:"listen"`          // The address to listen on.
	Host       string `json:"host"`            // The hostname to whitelist.
	Email      string `json:"email,omitempty"` // The email address for autocert.
}

// Load loads the certificates from the cache or the configuration.
func (c *TLSConfig) Load(certCache autocert.Cache) (*tls.Config, error) {
	if c.Host == "" {
		return nil, errors.New("unable to request a certificate, no host name configured")
	}

	// Default to disk cache
	if certCache == nil {
		certCache = autocert.DirCache("certs")
	}

	// Create an auto-cert manager
	certManager := autocert.Manager{
		Prompt:     autocert.AcceptTOS,
		HostPolicy: autocert.HostWhitelist(c.Host),
		Email:      c.Email,
		Cache:      certCache,
	}

	return &tls.Config{
		GetCertificate: certManager.GetCertificate,
	}, nil
}

// VaultConfig represents Vault configuration.
type VaultConfig struct {
	Address     string `json:"address"` // The vault address to use.
	Application string `json:"app"`     // The vault application ID to use.
}

// NewClient creates a new vault client for the configuration.
func (c *VaultConfig) NewClient(user string) (client *VaultClient, err error) {
	if c.Address == "" || c.Application == "" {
		return nil, errors.New("unable to configure Vault provider")
	}

	client = NewVaultClient(c.Address)
	err = client.Authenticate(c.Application, user)
	return
}

// ProviderConfig represents provider configuration.
type ProviderConfig struct {

	// The storage provider, this can either be specific builtin or the plugin path (file or
	// url) if the plugin is specified, it must contain a constructor function named 'New'
	// which returns an interface{}.
	Provider string `json:"provider"`

	// The configuration for a provider. This specifies various parameters to provide to the
	// specific provider during the Configure() call.
	Config map[string]interface{} `json:"config,omitempty"`
}

// LoadOrPanic loads a provider from the configuration and uses one or several builtins
// provided. If the provider is not found, it panics.
func (c *ProviderConfig) LoadOrPanic(builtins ...Provider) Provider {
	provider, err := c.Load(builtins...)
	if err != nil {
		panic(err)
	}

	return provider
}

// Load loads a provider from the configuration and uses one or several builtins provided.
func (c *ProviderConfig) Load(builtins ...Provider) (Provider, error) {
	for _, builtin := range builtins {
		if strings.ToLower(builtin.Name()) == strings.ToLower(c.Provider) {
			if err := builtin.Configure(c.Config); err != nil {
				return nil, errors.New("The provider '" + c.Provider + "' could not be loaded. " + err.Error())
			}

			return builtin, nil
		}
	}

	return c.LoadPlugin()
}

// LoadProvider loads a provider from the configuration or panics if the configuration is
// specified, but the provider was not found or not able to configure. This uses the first
// provider as a default value.
func LoadProvider(config *ProviderConfig, providers ...Provider) Provider {
	if config == nil || config.Provider == "" {
		return providers[0]
	}

	// Load the provider according to the configuration
	return config.LoadOrPanic(providers...)
}

// Write writes the configuration to a specific writer, in JSON format.
func write(config interface{}, output io.Writer) (int, error) {
	var formatted bytes.Buffer
	body, err := json.Marshal(config)
	if err != nil {
		return 0, err
	}

	if err := json.Indent(&formatted, body, "", "\t"); err != nil {
		return 0, err
	}

	return output.Write(formatted.Bytes())
}

// createDefault writes the default configuration to disk.
func createDefault(path string, newDefault func() Config) (Config, error) {
	f, err := os.OpenFile(path, os.O_CREATE|os.O_RDWR, os.ModePerm)
	if err != nil {
		return nil, err
	}

	defer f.Close()
	c := newDefault()
	if _, err := write(c, f); err != nil {
		return nil, err
	}
	if err := f.Sync(); err != nil {
		return nil, err
	}
	return c, nil
}

// ReadOrCreate reads or creates the configuration object.
func ReadOrCreate(prefix string, path string, newDefault func() Config, stores ...SecretStore) (cfg Config, err error) {
	cfg = newDefault()
	if _, err := os.Stat(path); os.IsNotExist(err) {
		// Create a configuration and write it to a file
		if cfg, err = createDefault(path, newDefault); err != nil {
			return nil, err
		}
	} else {
		// Read the config from file
		b, err := ioutil.ReadFile(path)
		if err != nil {
			return nil, err
		}

		// Unmarshal the configuration
		if err := json.Unmarshal(b, cfg); err != nil {
			return nil, err
		}
	}

	// Apply all the store overrides, in order
	for _, store := range stores {
		if err := store.Configure(cfg); err == nil {
			declassify(cfg, prefix, store)
		}
	}

	return cfg, nil
}

// Declassify traverses the configuration and resolves secrets.
func declassify(config interface{}, prefix string, provider SecretStore) {
	original := reflect.ValueOf(config)
	declassifyRecursive(prefix, provider, original)
}

// DeclassifyRecursive traverses the configuration and resolves secrets.
func declassifyRecursive(prefix string, provider SecretStore, value reflect.Value) {
	switch value.Kind() {
	case reflect.Ptr:
		pValue := value.Elem()
		if !pValue.IsValid() {
			// Create a new struct and set the value
			pValue = reflect.New(value.Type().Elem())
			value.Set(pValue)
		}

		declassifyRecursive(prefix, provider, pValue)

	// If it is a struct we translate each field
	case reflect.Struct:
		for i := 0; i < value.NumField(); i++ {
			name := getFieldName(value.Type().Field(i))
			declassifyRecursive(prefix+"/"+name, provider, value.Field(i))
		}

	// This is a integer, we need to fetch the secret
	case reflect.Int:
		if v, ok := provider.GetSecret(prefix); ok {
			if iv, err := strconv.ParseInt(v, 10, 64); err == nil {
				value.SetInt(iv)
			}
		}

	// This is a string, we need to fetch the secret
	case reflect.String:
		if v, ok := provider.GetSecret(prefix); ok {
			value.SetString(v)
		}

	// This is a map, unmarshal and set
	case reflect.Map:
		if v, ok := provider.GetSecret(prefix); ok {
			var out map[string]interface{}
			if err := json.Unmarshal([]byte(v), &out); err == nil {
				value.Set(reflect.ValueOf(out))
			}
		}
	}
}

func getFieldName(f reflect.StructField) string {
	return strings.Replace(string(f.Tag.Get("json")), ",omitempty", "", -1)
}

func resolvePath(path string) string {

	// If it's an url, download the file
	if strings.HasPrefix(path, "http") {
		f, err := httpFile(path)
		if err != nil {
			panic(err)
		}

		// Get the downloaded file path
		path = f.Name()
	}

	// Make sure the path is absolute
	path, _ = filepath.Abs(path)
	return path
}
<commit_msg>Support for internal Certification Autority<commit_after>package config

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"strconv"
	"strings"

	"golang.org/x/crypto/acme/autocert"
)

// SecretStore represents a contract for a store capable of resolving secrets.
type SecretStore interface {
	Configure(config Config) error
	GetSecret(secretName string) (string, bool)
}

// CertificateStore represents a secret store which can be used for certificates.
type CertificateStore interface {
	autocert.Cache
}

// Provider represents a configurable provider.
type Provider interface {
	Name() string
	Configure(config map[string]interface{}) error
}

// Config represents a configuration interface.
type Config interface {
	Vault() *VaultConfig
}

// TLSConfig represents TLS listener configuration.
type TLSConfig struct {
	ListenAddr  string `json:"listen"`                // The address to listen on.
	Host        string `json:"host"`                  // The hostname to whitelist.
	Email       string `json:"email,omitempty"`       // The email address for autocert.
	Certificate string `json:"certificate,omitempty"` // The certificate request.
	PrivateKey  string `json:"private,omitempty"`     // The private key for the certificate.
}

// Load loads the certificates from the cache or the configuration.
func (c *TLSConfig) Load(certCache autocert.Cache) (*tls.Config, error) {
	if c.Certificate != "" {
		if c.Certificate == "" || c.PrivateKey == "" {
			return &tls.Config{}, errors.New("No certificate or private key configured")
		}

		// If the certificate provided is in plain text, write to file so we can read it.
		if strings.HasPrefix(c.Certificate, "---") {
			if err := ioutil.WriteFile("broker.crt", []byte(c.Certificate), os.ModePerm); err == nil {
				c.Certificate = "broker.crt"
			}
		}

		// If the private key provided is in plain text, write to file so we can read it.
		if strings.HasPrefix(c.PrivateKey, "---") {
			if err := ioutil.WriteFile("broker.key", []byte(c.PrivateKey), os.ModePerm); err == nil {
				c.PrivateKey = "broker.key"
			}
		}

		// Make sure the paths are absolute, otherwise we won't be able to read the files.
		c.Certificate = resolvePath(c.Certificate)
		c.PrivateKey = resolvePath(c.PrivateKey)

		// Load the certificate from the cert/key files.

		cer, err := tls.LoadX509KeyPair(c.Certificate, c.PrivateKey)
		return &tls.Config{
			Certificates: []tls.Certificate{cer},
		}, err
	} else {

		if c.Host == "" {
			return nil, errors.New("unable to request a certificate, no host name configured")
		}

		// Default to disk cache
		if certCache == nil {
			certCache = autocert.DirCache("certs")
		}

		// Create an auto-cert manager
		certManager := autocert.Manager{
			Prompt:     autocert.AcceptTOS,
			HostPolicy: autocert.HostWhitelist(c.Host),
			Email:      c.Email,
			Cache:      certCache,
		}

		return &tls.Config{
			GetCertificate: certManager.GetCertificate,
		}, nil
	}

}

// VaultConfig represents Vault configuration.
type VaultConfig struct {
	Address     string `json:"address"` // The vault address to use.
	Application string `json:"app"`     // The vault application ID to use.
}

// NewClient creates a new vault client for the configuration.
func (c *VaultConfig) NewClient(user string) (client *VaultClient, err error) {
	if c.Address == "" || c.Application == "" {
		return nil, errors.New("unable to configure Vault provider")
	}

	client = NewVaultClient(c.Address)
	err = client.Authenticate(c.Application, user)
	return
}

// ProviderConfig represents provider configuration.
type ProviderConfig struct {

	// The storage provider, this can either be specific builtin or the plugin path (file or
	// url) if the plugin is specified, it must contain a constructor function named 'New'
	// which returns an interface{}.
	Provider string `json:"provider"`

	// The configuration for a provider. This specifies various parameters to provide to the
	// specific provider during the Configure() call.
	Config map[string]interface{} `json:"config,omitempty"`
}

// LoadOrPanic loads a provider from the configuration and uses one or several builtins
// provided. If the provider is not found, it panics.
func (c *ProviderConfig) LoadOrPanic(builtins ...Provider) Provider {
	provider, err := c.Load(builtins...)
	if err != nil {
		panic(err)
	}

	return provider
}

// Load loads a provider from the configuration and uses one or several builtins provided.
func (c *ProviderConfig) Load(builtins ...Provider) (Provider, error) {
	for _, builtin := range builtins {
		if strings.ToLower(builtin.Name()) == strings.ToLower(c.Provider) {
			if err := builtin.Configure(c.Config); err != nil {
				return nil, errors.New("The provider '" + c.Provider + "' could not be loaded. " + err.Error())
			}

			return builtin, nil
		}
	}

	return c.LoadPlugin()
}

// LoadProvider loads a provider from the configuration or panics if the configuration is
// specified, but the provider was not found or not able to configure. This uses the first
// provider as a default value.
func LoadProvider(config *ProviderConfig, providers ...Provider) Provider {
	if config == nil || config.Provider == "" {
		return providers[0]
	}

	// Load the provider according to the configuration
	return config.LoadOrPanic(providers...)
}

// Write writes the configuration to a specific writer, in JSON format.
func write(config interface{}, output io.Writer) (int, error) {
	var formatted bytes.Buffer
	body, err := json.Marshal(config)
	if err != nil {
		return 0, err
	}

	if err := json.Indent(&formatted, body, "", "\t"); err != nil {
		return 0, err
	}

	return output.Write(formatted.Bytes())
}

// createDefault writes the default configuration to disk.
func createDefault(path string, newDefault func() Config) (Config, error) {
	f, err := os.OpenFile(path, os.O_CREATE|os.O_RDWR, os.ModePerm)
	if err != nil {
		return nil, err
	}

	defer f.Close()
	c := newDefault()
	if _, err := write(c, f); err != nil {
		return nil, err
	}
	if err := f.Sync(); err != nil {
		return nil, err
	}
	return c, nil
}

// ReadOrCreate reads or creates the configuration object.
func ReadOrCreate(prefix string, path string, newDefault func() Config, stores ...SecretStore) (cfg Config, err error) {
	cfg = newDefault()
	if _, err := os.Stat(path); os.IsNotExist(err) {
		// Create a configuration and write it to a file
		if cfg, err = createDefault(path, newDefault); err != nil {
			return nil, err
		}
	} else {
		// Read the config from file
		b, err := ioutil.ReadFile(path)
		if err != nil {
			return nil, err
		}

		// Unmarshal the configuration
		if err := json.Unmarshal(b, cfg); err != nil {
			return nil, err
		}
	}

	// Apply all the store overrides, in order
	for _, store := range stores {
		if err := store.Configure(cfg); err == nil {
			declassify(cfg, prefix, store)
		}
	}

	return cfg, nil
}

// Declassify traverses the configuration and resolves secrets.
func declassify(config interface{}, prefix string, provider SecretStore) {
	original := reflect.ValueOf(config)
	declassifyRecursive(prefix, provider, original)
}

// DeclassifyRecursive traverses the configuration and resolves secrets.
func declassifyRecursive(prefix string, provider SecretStore, value reflect.Value) {
	switch value.Kind() {
	case reflect.Ptr:
		pValue := value.Elem()
		if !pValue.IsValid() {
			// Create a new struct and set the value
			pValue = reflect.New(value.Type().Elem())
			value.Set(pValue)
		}

		declassifyRecursive(prefix, provider, pValue)

	// If it is a struct we translate each field
	case reflect.Struct:
		for i := 0; i < value.NumField(); i++ {
			name := getFieldName(value.Type().Field(i))
			declassifyRecursive(prefix+"/"+name, provider, value.Field(i))
		}

	// This is a integer, we need to fetch the secret
	case reflect.Int:
		if v, ok := provider.GetSecret(prefix); ok {
			if iv, err := strconv.ParseInt(v, 10, 64); err == nil {
				value.SetInt(iv)
			}
		}

	// This is a string, we need to fetch the secret
	case reflect.String:
		if v, ok := provider.GetSecret(prefix); ok {
			value.SetString(v)
		}

	// This is a map, unmarshal and set
	case reflect.Map:
		if v, ok := provider.GetSecret(prefix); ok {
			var out map[string]interface{}
			if err := json.Unmarshal([]byte(v), &out); err == nil {
				value.Set(reflect.ValueOf(out))
			}
		}
	}
}

func getFieldName(f reflect.StructField) string {
	return strings.Replace(string(f.Tag.Get("json")), ",omitempty", "", -1)
}

func resolvePath(path string) string {

	// If it's an url, download the file
	if strings.HasPrefix(path, "http") {
		f, err := httpFile(path)
		if err != nil {
			panic(err)
		}

		// Get the downloaded file path
		path = f.Name()
	}

	// Make sure the path is absolute
	path, _ = filepath.Abs(path)
	return path
}
<|endoftext|>


<commit_before>package command

import (
	"context"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/filer"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/remote_pb"
	"github.com/chrislusf/seaweedfs/weed/remote_storage"
	"github.com/chrislusf/seaweedfs/weed/replication/source"
	"github.com/chrislusf/seaweedfs/weed/security"
	"github.com/chrislusf/seaweedfs/weed/util"
	"google.golang.org/grpc"
	"time"
)

type RemoteSyncOptions struct {
	filerAddress       *string
	grpcDialOption     grpc.DialOption
	readChunkFromFiler *bool
	debug              *bool
	timeAgo            *time.Duration
	dir                *string
}

const (
	RemoteSyncKeyPrefix = "remote.sync."
)

var _ = filer_pb.FilerClient(&RemoteSyncOptions{})

func (option *RemoteSyncOptions) WithFilerClient(fn func(filer_pb.SeaweedFilerClient) error) error {
	return pb.WithFilerClient(*option.filerAddress, option.grpcDialOption, func(client filer_pb.SeaweedFilerClient) error {
		return fn(client)
	})
}
func (option *RemoteSyncOptions) AdjustedUrl(location *filer_pb.Location) string {
	return location.Url
}

var (
	remoteSyncOptions RemoteSyncOptions
)

func init() {
	cmdFilerRemoteSynchronize.Run = runFilerRemoteSynchronize // break init cycle
	remoteSyncOptions.filerAddress = cmdFilerRemoteSynchronize.Flag.String("filer", "localhost:8888", "filer of the SeaweedFS cluster")
	remoteSyncOptions.dir = cmdFilerRemoteSynchronize.Flag.String("dir", "/", "a mounted directory on filer")
	remoteSyncOptions.readChunkFromFiler = cmdFilerRemoteSynchronize.Flag.Bool("filerProxy", false, "read file chunks from filer instead of volume servers")
	remoteSyncOptions.debug = cmdFilerRemoteSynchronize.Flag.Bool("debug", false, "debug mode to print out filer updated remote files")
	remoteSyncOptions.timeAgo = cmdFilerRemoteSynchronize.Flag.Duration("timeAgo", 0, "start time before now. \"300ms\", \"1.5h\" or \"2h45m\". Valid time units are \"ns\", \"us\" (or \"µs\"), \"ms\", \"s\", \"m\", \"h\"")
}

var cmdFilerRemoteSynchronize = &Command{
	UsageLine: "filer.remote.sync -filer=<filerHost>:<filerPort> -dir=/mount/s3_on_cloud",
	Short:     "resumable continuously write back updates to remote storage if the directory is mounted to the remote storage",
	Long: `resumable continuously write back updates to remote storage if the directory is mounted to the remote storage

	filer.remote.sync listens on filer update events.
	If any mounted remote file is updated, it will fetch the updated content,
	and write to the remote storage.
`,
}

func runFilerRemoteSynchronize(cmd *Command, args []string) bool {

	util.LoadConfiguration("security", false)
	grpcDialOption := security.LoadClientTLS(util.GetViper(), "grpc.client")
	remoteSyncOptions.grpcDialOption = grpcDialOption

	dir := *remoteSyncOptions.dir
	filerAddress := *remoteSyncOptions.filerAddress

	// read filer remote storage mount mappings
	_, _, remoteStorageMountLocation, storageConf, detectErr := filer.DetectMountInfo(grpcDialOption, filerAddress, dir)
	if detectErr != nil {
		fmt.Printf("read mount info: %v", detectErr)
		return false
	}

	filerSource := &source.FilerSource{}
	filerSource.DoInitialize(
		filerAddress,
		pb.ServerToGrpcAddress(filerAddress),
		"/", // does not matter
		*remoteSyncOptions.readChunkFromFiler,
	)

	fmt.Printf("synchronize %s to remote storage...\n", dir)
	util.RetryForever("filer.remote.sync "+dir, func() error {
		return followUpdatesAndUploadToRemote(&remoteSyncOptions, filerSource, dir, storageConf, remoteStorageMountLocation)
	}, func(err error) bool {
		if err != nil {
			glog.Errorf("synchronize %s: %v", dir, err)
		}
		return true
	})

	return true
}

func followUpdatesAndUploadToRemote(option *RemoteSyncOptions, filerSource *source.FilerSource, mountedDir string, remoteStorage *remote_pb.RemoteConf, remoteStorageMountLocation *remote_pb.RemoteStorageLocation) error {

	dirHash := util.HashStringToLong(mountedDir)

	// 1. specified by timeAgo
	// 2. last offset timestamp for this directory
	// 3. directory creation time
	var lastOffsetTs time.Time
	if *option.timeAgo == 0 {
		mountedDirEntry, err := filer_pb.GetEntry(option, util.FullPath(mountedDir))
		if err != nil {
			return fmt.Errorf("lookup %s: %v", mountedDir, err)
		}

		lastOffsetTsNs, err := getOffset(option.grpcDialOption, *option.filerAddress, RemoteSyncKeyPrefix, int32(dirHash))
		if mountedDirEntry != nil {
			if err == nil && mountedDirEntry.Attributes.Crtime < lastOffsetTsNs/1000000 {
				lastOffsetTs = time.Unix(0, lastOffsetTsNs)
				glog.V(0).Infof("resume from %v", lastOffsetTs)
			} else {
				lastOffsetTs = time.Unix(mountedDirEntry.Attributes.Crtime, 0)
			}
		} else {
			lastOffsetTs = time.Now()
		}
	} else {
		lastOffsetTs = time.Now().Add(-*option.timeAgo)
	}

	client, err := remote_storage.GetRemoteStorage(remoteStorage)
	if err != nil {
		return err
	}

	eachEntryFunc := func(resp *filer_pb.SubscribeMetadataResponse) error {
		message := resp.EventNotification
		if message.OldEntry == nil && message.NewEntry == nil {
			return nil
		}
		if message.OldEntry == nil && message.NewEntry != nil {
			if !filer.HasData(message.NewEntry) {
				return nil
			}
			glog.V(2).Infof("create: %+v", resp)
			if !shouldSendToRemote(message.NewEntry) {
				glog.V(2).Infof("skipping creating: %+v", resp)
				return nil
			}
			dest := toRemoteStorageLocation(util.FullPath(mountedDir), util.NewFullPath(message.NewParentPath, message.NewEntry.Name), remoteStorageMountLocation)
			if message.NewEntry.IsDirectory {
				glog.V(0).Infof("mkdir  %s", remote_storage.FormatLocation(dest))
				return client.WriteDirectory(dest, message.NewEntry)
			}
			glog.V(0).Infof("create %s", remote_storage.FormatLocation(dest))
			reader := filer.NewFileReader(filerSource, message.NewEntry)
			remoteEntry, writeErr := client.WriteFile(dest, message.NewEntry, reader)
			if writeErr != nil {
				return writeErr
			}
			return updateLocalEntry(&remoteSyncOptions, message.NewParentPath, message.NewEntry, remoteEntry)
		}
		if message.OldEntry != nil && message.NewEntry == nil {
			glog.V(2).Infof("delete: %+v", resp)
			dest := toRemoteStorageLocation(util.FullPath(mountedDir), util.NewFullPath(resp.Directory, message.OldEntry.Name), remoteStorageMountLocation)
			glog.V(0).Infof("delete %s", remote_storage.FormatLocation(dest))
			return client.DeleteFile(dest)
		}
		if message.OldEntry != nil && message.NewEntry != nil {
			oldDest := toRemoteStorageLocation(util.FullPath(mountedDir), util.NewFullPath(resp.Directory, message.OldEntry.Name), remoteStorageMountLocation)
			dest := toRemoteStorageLocation(util.FullPath(mountedDir), util.NewFullPath(message.NewParentPath, message.NewEntry.Name), remoteStorageMountLocation)
			if !shouldSendToRemote(message.NewEntry) {
				glog.V(2).Infof("skipping updating: %+v", resp)
				return nil
			}
			if message.NewEntry.IsDirectory {
				return client.WriteDirectory(dest, message.NewEntry)
			}
			if resp.Directory == message.NewParentPath && message.OldEntry.Name == message.NewEntry.Name {
				if filer.IsSameData(message.OldEntry, message.NewEntry) {
					glog.V(2).Infof("update meta: %+v", resp)
					return client.UpdateFileMetadata(dest, message.OldEntry, message.NewEntry)
				}
			}
			glog.V(2).Infof("update: %+v", resp)
			glog.V(0).Infof("delete %s", remote_storage.FormatLocation(oldDest))
			if err := client.DeleteFile(oldDest); err != nil {
				return err
			}
			reader := filer.NewFileReader(filerSource, message.NewEntry)
			glog.V(0).Infof("create %s", remote_storage.FormatLocation(dest))
			remoteEntry, writeErr := client.WriteFile(dest, message.NewEntry, reader)
			if writeErr != nil {
				return writeErr
			}
			return updateLocalEntry(&remoteSyncOptions, message.NewParentPath, message.NewEntry, remoteEntry)
		}

		return nil
	}

	processEventFnWithOffset := pb.AddOffsetFunc(eachEntryFunc, 3*time.Second, func(counter int64, lastTsNs int64) error {
		lastTime := time.Unix(0, lastTsNs)
		glog.V(0).Infof("remote sync %s progressed to %v %0.2f/sec", *option.filerAddress, lastTime, float64(counter)/float64(3))
		return setOffset(option.grpcDialOption, *option.filerAddress, RemoteSyncKeyPrefix, int32(dirHash), lastTsNs)
	})

	return pb.FollowMetadata(*option.filerAddress, option.grpcDialOption,
		"filer.remote.sync", mountedDir, lastOffsetTs.UnixNano(), 0, processEventFnWithOffset, false)
}

func toRemoteStorageLocation(mountDir, sourcePath util.FullPath, remoteMountLocation *remote_pb.RemoteStorageLocation) *remote_pb.RemoteStorageLocation {
	source := string(sourcePath[len(mountDir):])
	dest := util.FullPath(remoteMountLocation.Path).Child(source)
	return &remote_pb.RemoteStorageLocation{
		Name:   remoteMountLocation.Name,
		Bucket: remoteMountLocation.Bucket,
		Path:   string(dest),
	}
}

func shouldSendToRemote(entry *filer_pb.Entry) bool {
	if entry.RemoteEntry == nil {
		return true
	}
	if entry.RemoteEntry.LastLocalSyncTsNs/1e9 < entry.Attributes.Mtime {
		return true
	}
	return false
}

func updateLocalEntry(filerClient filer_pb.FilerClient, dir string, entry *filer_pb.Entry, remoteEntry *filer_pb.RemoteEntry) error {
	entry.RemoteEntry = remoteEntry
	return filerClient.WithFilerClient(func(client filer_pb.SeaweedFilerClient) error {
		_, err := client.UpdateEntry(context.Background(), &filer_pb.UpdateEntryRequest{
			Directory: dir,
			Entry:     entry,
		})
		return err
	})
}
<commit_msg>refactoring<commit_after>package command

import (
	"context"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/filer"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/remote_pb"
	"github.com/chrislusf/seaweedfs/weed/remote_storage"
	"github.com/chrislusf/seaweedfs/weed/replication/source"
	"github.com/chrislusf/seaweedfs/weed/security"
	"github.com/chrislusf/seaweedfs/weed/util"
	"google.golang.org/grpc"
	"time"
)

type RemoteSyncOptions struct {
	filerAddress       *string
	grpcDialOption     grpc.DialOption
	readChunkFromFiler *bool
	debug              *bool
	timeAgo            *time.Duration
	dir                *string
}

const (
	RemoteSyncKeyPrefix = "remote.sync."
)

var _ = filer_pb.FilerClient(&RemoteSyncOptions{})

func (option *RemoteSyncOptions) WithFilerClient(fn func(filer_pb.SeaweedFilerClient) error) error {
	return pb.WithFilerClient(*option.filerAddress, option.grpcDialOption, func(client filer_pb.SeaweedFilerClient) error {
		return fn(client)
	})
}
func (option *RemoteSyncOptions) AdjustedUrl(location *filer_pb.Location) string {
	return location.Url
}

var (
	remoteSyncOptions RemoteSyncOptions
)

func init() {
	cmdFilerRemoteSynchronize.Run = runFilerRemoteSynchronize // break init cycle
	remoteSyncOptions.filerAddress = cmdFilerRemoteSynchronize.Flag.String("filer", "localhost:8888", "filer of the SeaweedFS cluster")
	remoteSyncOptions.dir = cmdFilerRemoteSynchronize.Flag.String("dir", "/", "a mounted directory on filer")
	remoteSyncOptions.readChunkFromFiler = cmdFilerRemoteSynchronize.Flag.Bool("filerProxy", false, "read file chunks from filer instead of volume servers")
	remoteSyncOptions.debug = cmdFilerRemoteSynchronize.Flag.Bool("debug", false, "debug mode to print out filer updated remote files")
	remoteSyncOptions.timeAgo = cmdFilerRemoteSynchronize.Flag.Duration("timeAgo", 0, "start time before now. \"300ms\", \"1.5h\" or \"2h45m\". Valid time units are \"ns\", \"us\" (or \"µs\"), \"ms\", \"s\", \"m\", \"h\"")
}

var cmdFilerRemoteSynchronize = &Command{
	UsageLine: "filer.remote.sync -filer=<filerHost>:<filerPort> -dir=/mount/s3_on_cloud",
	Short:     "resumable continuously write back updates to remote storage if the directory is mounted to the remote storage",
	Long: `resumable continuously write back updates to remote storage if the directory is mounted to the remote storage

	filer.remote.sync listens on filer update events.
	If any mounted remote file is updated, it will fetch the updated content,
	and write to the remote storage.
`,
}

func runFilerRemoteSynchronize(cmd *Command, args []string) bool {

	util.LoadConfiguration("security", false)
	grpcDialOption := security.LoadClientTLS(util.GetViper(), "grpc.client")
	remoteSyncOptions.grpcDialOption = grpcDialOption

	dir := *remoteSyncOptions.dir
	filerAddress := *remoteSyncOptions.filerAddress

	filerSource := &source.FilerSource{}
	filerSource.DoInitialize(
		filerAddress,
		pb.ServerToGrpcAddress(filerAddress),
		"/", // does not matter
		*remoteSyncOptions.readChunkFromFiler,
	)

	fmt.Printf("synchronize %s to remote storage...\n", dir)
	util.RetryForever("filer.remote.sync "+dir, func() error {
		return followUpdatesAndUploadToRemote(&remoteSyncOptions, filerSource, dir)
	}, func(err error) bool {
		if err != nil {
			glog.Errorf("synchronize %s: %v", dir, err)
		}
		return true
	})

	return true
}

func followUpdatesAndUploadToRemote(option *RemoteSyncOptions, filerSource *source.FilerSource, mountedDir string) error {

	// read filer remote storage mount mappings
	_, _, remoteStorageMountLocation, remoteStorage, detectErr := filer.DetectMountInfo(option.grpcDialOption, *option.filerAddress, mountedDir)
	if detectErr != nil {
		return fmt.Errorf("read mount info: %v", detectErr)
	}

	dirHash := util.HashStringToLong(mountedDir)

	// 1. specified by timeAgo
	// 2. last offset timestamp for this directory
	// 3. directory creation time
	var lastOffsetTs time.Time
	if *option.timeAgo == 0 {
		mountedDirEntry, err := filer_pb.GetEntry(option, util.FullPath(mountedDir))
		if err != nil {
			return fmt.Errorf("lookup %s: %v", mountedDir, err)
		}

		lastOffsetTsNs, err := getOffset(option.grpcDialOption, *option.filerAddress, RemoteSyncKeyPrefix, int32(dirHash))
		if mountedDirEntry != nil {
			if err == nil && mountedDirEntry.Attributes.Crtime < lastOffsetTsNs/1000000 {
				lastOffsetTs = time.Unix(0, lastOffsetTsNs)
				glog.V(0).Infof("resume from %v", lastOffsetTs)
			} else {
				lastOffsetTs = time.Unix(mountedDirEntry.Attributes.Crtime, 0)
			}
		} else {
			lastOffsetTs = time.Now()
		}
	} else {
		lastOffsetTs = time.Now().Add(-*option.timeAgo)
	}

	client, err := remote_storage.GetRemoteStorage(remoteStorage)
	if err != nil {
		return err
	}

	eachEntryFunc := func(resp *filer_pb.SubscribeMetadataResponse) error {
		message := resp.EventNotification
		if message.OldEntry == nil && message.NewEntry == nil {
			return nil
		}
		if message.OldEntry == nil && message.NewEntry != nil {
			if !filer.HasData(message.NewEntry) {
				return nil
			}
			glog.V(2).Infof("create: %+v", resp)
			if !shouldSendToRemote(message.NewEntry) {
				glog.V(2).Infof("skipping creating: %+v", resp)
				return nil
			}
			dest := toRemoteStorageLocation(util.FullPath(mountedDir), util.NewFullPath(message.NewParentPath, message.NewEntry.Name), remoteStorageMountLocation)
			if message.NewEntry.IsDirectory {
				glog.V(0).Infof("mkdir  %s", remote_storage.FormatLocation(dest))
				return client.WriteDirectory(dest, message.NewEntry)
			}
			glog.V(0).Infof("create %s", remote_storage.FormatLocation(dest))
			reader := filer.NewFileReader(filerSource, message.NewEntry)
			remoteEntry, writeErr := client.WriteFile(dest, message.NewEntry, reader)
			if writeErr != nil {
				return writeErr
			}
			return updateLocalEntry(&remoteSyncOptions, message.NewParentPath, message.NewEntry, remoteEntry)
		}
		if message.OldEntry != nil && message.NewEntry == nil {
			glog.V(2).Infof("delete: %+v", resp)
			dest := toRemoteStorageLocation(util.FullPath(mountedDir), util.NewFullPath(resp.Directory, message.OldEntry.Name), remoteStorageMountLocation)
			glog.V(0).Infof("delete %s", remote_storage.FormatLocation(dest))
			return client.DeleteFile(dest)
		}
		if message.OldEntry != nil && message.NewEntry != nil {
			oldDest := toRemoteStorageLocation(util.FullPath(mountedDir), util.NewFullPath(resp.Directory, message.OldEntry.Name), remoteStorageMountLocation)
			dest := toRemoteStorageLocation(util.FullPath(mountedDir), util.NewFullPath(message.NewParentPath, message.NewEntry.Name), remoteStorageMountLocation)
			if !shouldSendToRemote(message.NewEntry) {
				glog.V(2).Infof("skipping updating: %+v", resp)
				return nil
			}
			if message.NewEntry.IsDirectory {
				return client.WriteDirectory(dest, message.NewEntry)
			}
			if resp.Directory == message.NewParentPath && message.OldEntry.Name == message.NewEntry.Name {
				if filer.IsSameData(message.OldEntry, message.NewEntry) {
					glog.V(2).Infof("update meta: %+v", resp)
					return client.UpdateFileMetadata(dest, message.OldEntry, message.NewEntry)
				}
			}
			glog.V(2).Infof("update: %+v", resp)
			glog.V(0).Infof("delete %s", remote_storage.FormatLocation(oldDest))
			if err := client.DeleteFile(oldDest); err != nil {
				return err
			}
			reader := filer.NewFileReader(filerSource, message.NewEntry)
			glog.V(0).Infof("create %s", remote_storage.FormatLocation(dest))
			remoteEntry, writeErr := client.WriteFile(dest, message.NewEntry, reader)
			if writeErr != nil {
				return writeErr
			}
			return updateLocalEntry(&remoteSyncOptions, message.NewParentPath, message.NewEntry, remoteEntry)
		}

		return nil
	}

	processEventFnWithOffset := pb.AddOffsetFunc(eachEntryFunc, 3*time.Second, func(counter int64, lastTsNs int64) error {
		lastTime := time.Unix(0, lastTsNs)
		glog.V(0).Infof("remote sync %s progressed to %v %0.2f/sec", *option.filerAddress, lastTime, float64(counter)/float64(3))
		return setOffset(option.grpcDialOption, *option.filerAddress, RemoteSyncKeyPrefix, int32(dirHash), lastTsNs)
	})

	return pb.FollowMetadata(*option.filerAddress, option.grpcDialOption,
		"filer.remote.sync", mountedDir, lastOffsetTs.UnixNano(), 0, processEventFnWithOffset, false)
}

func toRemoteStorageLocation(mountDir, sourcePath util.FullPath, remoteMountLocation *remote_pb.RemoteStorageLocation) *remote_pb.RemoteStorageLocation {
	source := string(sourcePath[len(mountDir):])
	dest := util.FullPath(remoteMountLocation.Path).Child(source)
	return &remote_pb.RemoteStorageLocation{
		Name:   remoteMountLocation.Name,
		Bucket: remoteMountLocation.Bucket,
		Path:   string(dest),
	}
}

func shouldSendToRemote(entry *filer_pb.Entry) bool {
	if entry.RemoteEntry == nil {
		return true
	}
	if entry.RemoteEntry.LastLocalSyncTsNs/1e9 < entry.Attributes.Mtime {
		return true
	}
	return false
}

func updateLocalEntry(filerClient filer_pb.FilerClient, dir string, entry *filer_pb.Entry, remoteEntry *filer_pb.RemoteEntry) error {
	entry.RemoteEntry = remoteEntry
	return filerClient.WithFilerClient(func(client filer_pb.SeaweedFilerClient) error {
		_, err := client.UpdateEntry(context.Background(), &filer_pb.UpdateEntryRequest{
			Directory: dir,
			Entry:     entry,
		})
		return err
	})
}
<|endoftext|>