<commit_before>// Package adminz provides a simple set of adminz pages for administering
// a simple go server.
package adminz

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"runtime"
	"sync"
	"time"
)

type Adminz struct {
	// keep track of killfile state
	running bool

	// ticker that checks killfiles every 1 second
	killfileTicker *time.Ticker

	// list of killfilePaths to check
	killfilePaths []string

	// defaults to 1 second
	checkInterval time.Duration

	// generates data to return to /servicez endpoint. marshalled into json.
	servicez func() interface{}

	// resume is called when the server is unkilled
	onresume func()

	// pause is called when the server is killed
	onpause func()

	// healthy returns true iff the server is ready to respond to requests
	healthy func() bool

	sync.Mutex
}

// Creates a new Adminz "builder". Not safe to use until Build() is called.
func New() *Adminz {
	return &Adminz{}
}

func (a *Adminz) Resume() {
	a.Lock()
	defer a.Unlock()
	if !a.running {
		if a.onresume != nil {
			a.onresume()
		}
		a.running = true
	}
}

// Resume is called when the server is unkilled
func (a *Adminz) OnResume(resume func()) *Adminz {
	a.onresume = resume
	return a
}

// Pause the server if it is running and return true, otherwise no-op and return false.
func (a *Adminz) Pause() bool {
	a.Lock()
	defer a.Unlock()
	was := a.running
	if a.running {
		if a.onpause != nil {
			a.onpause()
		}
		a.running = false
	}
	return was
}

// pause is called when the server is killed
func (a *Adminz) OnPause(pause func()) *Adminz {
	a.onpause = pause
	return a
}

// healthy returns true iff the server is ready to respond to requests
func (a *Adminz) Healthy(healthy func() bool) *Adminz {
	a.healthy = healthy
	return a
}

// servicez generates data to return to /servicez endpoint. marshalled into
// json.
func (a *Adminz) Servicez(servicez func() interface{}) *Adminz {
	a.servicez = servicez
	return a
}

// Sets the list of killfilePaths to check.
func (a *Adminz) KillfilePaths(killfilePaths []string) *Adminz {
	a.killfilePaths = killfilePaths
	return a
}

// Sets frequency the killfile is checked. defaults every second
func (a *Adminz) KillfileInterval(interval time.Duration) *Adminz {
	a.checkInterval = interval
	return a
}

// Build initializes handlers and starts killfile checking. Make sure to
// remember to call this!
func (a *Adminz) Build() *Adminz {
	http.HandleFunc("/healthz", a.healthzHandler)
	http.HandleFunc("/servicez", a.servicezHandler)
	http.HandleFunc("/gc", a.gcHandler)

	log.Print("adminz registered")

	if a.checkInterval == 0 {
		a.checkInterval = 1 * time.Second
	}
	a.killfileTicker = time.NewTicker(a.checkInterval)

	// start killfile checking loop
	if len(a.killfilePaths) > 0 {
		log.Print("Watching paths for killfile: ", a.killfilePaths)
		a.checkKillfiles()
		go a.killfileLoop()
	} else {
		log.Print("No killfiles being watched.")
	}

	return a
}

func (a *Adminz) Stop() {
	if a.killfileTicker != nil {
		a.killfileTicker.Stop()
	}
}

// Generates the standard set of killfiles. Pass these to KillfilePaths
func Killfiles(ports ...int) []string {
	// the number of ports + the "all" killfile
	log.Print(ports)
	var ret = make([]string, len(ports)+1)
	for i, port := range ports {
		ret[i] = fmt.Sprintf("/dev/shm/healthz/kill.%d", port)
	}
	ret[len(ports)] = "/dev/shm/healthz/kill.all"
	return ret
}

func (a *Adminz) checkKillfiles() bool {
	for _, killfile := range a.killfilePaths {
		file, err := os.Open(killfile)
		if file != nil && err == nil {
			file.Close()
			if a.Pause() {
				log.Println("paused due to ", killfile)
			}
			return true
		}
	}
	a.Resume()
	return false
}

func (a *Adminz) killfileLoop() {
	for _ = range a.killfileTicker.C {
		a.checkKillfiles()
	}
}

func (a *Adminz) healthzHandler(w http.ResponseWriter, r *http.Request) {
	// we are healthy iff:
	// we are not killed AND
	// a.healthy is unset (so we ignore it) OR
	// a.healthy() returns true
	var ret string

	if !a.running && (a.healthy == nil || a.healthy()) {
		ret = "OK"
	} else {
		w.WriteHeader(http.StatusServiceUnavailable)
		ret = "Service Unavailable"
	}

	log.Print("Healthz returning ", ret)
	w.Write(([]byte)(ret))
}

func (a *Adminz) servicezHandler(w http.ResponseWriter, r *http.Request) {
	if a.servicez == nil {
		return
	}

	bytes, err := json.Marshal(a.servicez())
	if err == nil {
		w.Header().Add("Content-Type", "text/json")
		// TODO I probably need to serialize reads to servicez as who knows what
		// people will put in that function
		w.Write(bytes)
	} else {
		http.Error(w, err.Error(), 500)
	}
}

func (a *Adminz) gcHandler(w http.ResponseWriter, r *http.Request) {
	var mem runtime.MemStats

	mb := uint64(1024 * 1024)
	runtime.ReadMemStats(&mem)
	fmt.Fprintln(w, "Before")
	fmt.Fprintln(w, "\tAlloc\t", mem.Alloc/mb)
	fmt.Fprintln(w, "\tTotalAlloc:\t", mem.TotalAlloc/mb)
	fmt.Fprintln(w, "\tHeapAlloc:\t", mem.HeapAlloc/mb)
	fmt.Fprintln(w, "\tHeapSys:\t", mem.HeapSys/mb)
	fmt.Fprintln(w, "\tSys:\t", mem.Sys/mb)

	runtime.GC()

	runtime.ReadMemStats(&mem)
	fmt.Fprintln(w, "After")
	fmt.Fprintln(w, "\tAlloc\t", mem.Alloc/mb)
	fmt.Fprintln(w, "\tTotalAlloc:\t", mem.TotalAlloc/mb)
	fmt.Fprintln(w, "\tHeapAlloc:\t", mem.HeapAlloc/mb)
	fmt.Fprintln(w, "\tHeapSys:\t", mem.HeapSys/mb)
	fmt.Fprintln(w, "\tSys:\t", mem.Sys/mb)

	w.Write([]byte("OK"))
}
<commit_msg>pause/resume when GCing<commit_after>// Package adminz provides a simple set of adminz pages for administering
// a simple go server.
package adminz

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"runtime"
	"sync"
	"time"
)

type Adminz struct {
	// keep track of killfile state
	running bool

	// ticker that checks killfiles every 1 second
	killfileTicker *time.Ticker

	// list of killfilePaths to check
	killfilePaths []string

	// defaults to 1 second
	checkInterval time.Duration

	// generates data to return to /servicez endpoint. marshalled into json.
	servicez func() interface{}

	// resume is called when the server is unkilled
	onresume func()

	// pause is called when the server is killed
	onpause func()

	// healthy returns true iff the server is ready to respond to requests
	healthy func() bool

	sync.Mutex
}

// Creates a new Adminz "builder". Not safe to use until Build() is called.
func New() *Adminz {
	return &Adminz{}
}

func (a *Adminz) Resume() {
	a.Lock()
	defer a.Unlock()
	a.doResume()
}

// Internal helper to resume if stopped. MUST be called while holding a.Lock. Use a.Resume if not.
func (a *Adminz) doResume() {
	if !a.running {
		if a.onresume != nil {
			a.onresume()
		}
		a.running = true
	}
}

// Resume is called when the server is unkilled
func (a *Adminz) OnResume(resume func()) *Adminz {
	a.onresume = resume
	return a
}

// Pause the server if it is running and return true, otherwise no-op and return false.
func (a *Adminz) Pause() bool {
	a.Lock()
	defer a.Unlock()
	return a.doPause()
}

// Internal helper to pause if running. MUST be called while holding a.Lock. Use a.Pause if not.
func (a *Adminz) doPause() bool {
	was := a.running
	if a.running {
		if a.onpause != nil {
			a.onpause()
		}
		a.running = false
	}
	return was
}

// pause is called when the server is killed
func (a *Adminz) OnPause(pause func()) *Adminz {
	a.onpause = pause
	return a
}

// healthy returns true iff the server is ready to respond to requests
func (a *Adminz) Healthy(healthy func() bool) *Adminz {
	a.healthy = healthy
	return a
}

// servicez generates data to return to /servicez endpoint. marshalled into
// json.
func (a *Adminz) Servicez(servicez func() interface{}) *Adminz {
	a.servicez = servicez
	return a
}

// Sets the list of killfilePaths to check.
func (a *Adminz) KillfilePaths(killfilePaths []string) *Adminz {
	a.killfilePaths = killfilePaths
	return a
}

// Sets frequency the killfile is checked. defaults every second
func (a *Adminz) KillfileInterval(interval time.Duration) *Adminz {
	a.checkInterval = interval
	return a
}

// Build initializes handlers and starts killfile checking. Make sure to
// remember to call this!
func (a *Adminz) Build() *Adminz {
	http.HandleFunc("/healthz", a.healthzHandler)
	http.HandleFunc("/servicez", a.servicezHandler)
	http.HandleFunc("/gc", a.gcHandler)

	log.Print("adminz registered")

	if a.checkInterval == 0 {
		a.checkInterval = 1 * time.Second
	}
	a.killfileTicker = time.NewTicker(a.checkInterval)

	// start killfile checking loop
	if len(a.killfilePaths) > 0 {
		log.Print("Watching paths for killfile: ", a.killfilePaths)
		a.checkKillfiles()
		go a.killfileLoop()
	} else {
		log.Print("No killfiles being watched.")
	}

	return a
}

func (a *Adminz) Stop() {
	if a.killfileTicker != nil {
		a.killfileTicker.Stop()
	}
}

// Generates the standard set of killfiles. Pass these to KillfilePaths
func Killfiles(ports ...int) []string {
	// the number of ports + the "all" killfile
	log.Print(ports)
	var ret = make([]string, len(ports)+1)
	for i, port := range ports {
		ret[i] = fmt.Sprintf("/dev/shm/healthz/kill.%d", port)
	}
	ret[len(ports)] = "/dev/shm/healthz/kill.all"
	return ret
}

func (a *Adminz) checkKillfiles() bool {
	for _, killfile := range a.killfilePaths {
		file, err := os.Open(killfile)
		if file != nil && err == nil {
			file.Close()
			if a.Pause() {
				log.Println("paused due to ", killfile)
			}
			return true
		}
	}
	a.Resume()
	return false
}

func (a *Adminz) killfileLoop() {
	for _ = range a.killfileTicker.C {
		a.checkKillfiles()
	}
}

func (a *Adminz) healthzHandler(w http.ResponseWriter, r *http.Request) {
	// we are healthy iff:
	// we are not killed AND
	// a.healthy is unset (so we ignore it) OR
	// a.healthy() returns true
	var ret string

	if !a.running && (a.healthy == nil || a.healthy()) {
		ret = "OK"
	} else {
		w.WriteHeader(http.StatusServiceUnavailable)
		ret = "Service Unavailable"
	}

	log.Print("Healthz returning ", ret)
	w.Write(([]byte)(ret))
}

func (a *Adminz) servicezHandler(w http.ResponseWriter, r *http.Request) {
	if a.servicez == nil {
		return
	}

	bytes, err := json.Marshal(a.servicez())
	if err == nil {
		w.Header().Add("Content-Type", "text/json")
		// TODO I probably need to serialize reads to servicez as who knows what
		// people will put in that function
		w.Write(bytes)
	} else {
		http.Error(w, err.Error(), 500)
	}
}

func (a *Adminz) gcHandler(w http.ResponseWriter, r *http.Request) {
	var mem runtime.MemStats

	mb := uint64(1024 * 1024)
	runtime.ReadMemStats(&mem)
	fmt.Fprintln(w, "Before")
	fmt.Fprintln(w, "\tAlloc\t", mem.Alloc/mb)
	fmt.Fprintln(w, "\tTotalAlloc:\t", mem.TotalAlloc/mb)
	fmt.Fprintln(w, "\tHeapAlloc:\t", mem.HeapAlloc/mb)
	fmt.Fprintln(w, "\tHeapSys:\t", mem.HeapSys/mb)
	fmt.Fprintln(w, "\tSys:\t", mem.Sys/mb)

	a.Lock()
	was := a.doPause()
	runtime.GC()
	if was {
		a.doResume()
	}
	a.Unlock()

	runtime.ReadMemStats(&mem)
	fmt.Fprintln(w, "After")
	fmt.Fprintln(w, "\tAlloc\t", mem.Alloc/mb)
	fmt.Fprintln(w, "\tTotalAlloc:\t", mem.TotalAlloc/mb)
	fmt.Fprintln(w, "\tHeapAlloc:\t", mem.HeapAlloc/mb)
	fmt.Fprintln(w, "\tHeapSys:\t", mem.HeapSys/mb)
	fmt.Fprintln(w, "\tSys:\t", mem.Sys/mb)

	w.Write([]byte("OK"))
}
<|endoftext|>


<commit_before>package configurer

import (
	"os"
	"testing"

	"io/ioutil"

	"github.com/bpicode/fritzctl/config"
	"github.com/stretchr/testify/assert"
)

// TestInit test the intialization phase of the interactive cli.
func TestInit(t *testing.T) {
	cli := New().(*cliConfigurer)
	cli.ApplyDefaults(Defaults())
	assert.Equal(t, cli.defaultValues, cli.userValues)
	cli.userValues.file = "/tmp/throwaway.json"
	assert.NotEqual(t, cli.defaultValues, cli.userValues)
}

// TestObtain test the user data acquisition phase of the cli.
func TestObtain(t *testing.T) {
	cli := New().(*cliConfigurer)
	cli.ApplyDefaults(Defaults())
	exfg := cli.Obtain()
	assert.NotNil(t, exfg)
}

// TestWrite test the configuration write phase of the cli.
func TestWrite(t *testing.T) {
	cli := New().(*cliConfigurer)
	extendedCfg := Defaults()
	tf, _ := ioutil.TempFile("", "test_fritzctl.json.")
	defer tf.Close()
	defer os.Remove(tf.Name())
	extendedCfg.file = tf.Name()
	cli.ApplyDefaults(extendedCfg)
	err := cli.Write()
	assert.NoError(t, err)
}

// TestWrite test the configuration write phase of the cli.
func TestWriteAndRead(t *testing.T) {
	cli := New().(*cliConfigurer)
	extendedCfg := Defaults()
	tf, _ := ioutil.TempFile("", "test_fritzctl.json.")
	defer tf.Close()
	defer os.Remove(tf.Name())
	extendedCfg.file = tf.Name()
	cli.ApplyDefaults(extendedCfg)
	err := cli.Write()
	assert.NoError(t, err)
	re, err := config.New(tf.Name())
	assert.NoError(t, err)
	assert.NotNil(t, re)
	assert.Equal(t, *cli.userValues.fritzCfg.Net, *re.Net)
	assert.Equal(t, *cli.userValues.fritzCfg.Login, *re.Login)
	assert.Equal(t, *cli.userValues.fritzCfg.Pki, *re.Pki)
}

// TestWriteWithIOError test the write phase of the cli with error.
func TestWriteWithIOError(t *testing.T) {
	cli := New().(*cliConfigurer)
	extendedCfg := Defaults()
	extendedCfg.file = "/root/a/b/c/no/such/file/or/directory/cfg.json"
	cli.ApplyDefaults(extendedCfg)
	err := cli.Write()
	assert.Error(t, err)
}
<commit_msg>fix godoc<commit_after>package configurer

import (
	"os"
	"testing"

	"io/ioutil"

	"github.com/bpicode/fritzctl/config"
	"github.com/stretchr/testify/assert"
)

// TestInit test the intialization phase of the interactive cli.
func TestInit(t *testing.T) {
	cli := New().(*cliConfigurer)
	cli.ApplyDefaults(Defaults())
	assert.Equal(t, cli.defaultValues, cli.userValues)
	cli.userValues.file = "/tmp/throwaway.json"
	assert.NotEqual(t, cli.defaultValues, cli.userValues)
}

// TestObtain test the user data acquisition phase of the cli.
func TestObtain(t *testing.T) {
	cli := New().(*cliConfigurer)
	cli.ApplyDefaults(Defaults())
	exfg := cli.Obtain()
	assert.NotNil(t, exfg)
}

// TestWrite test the configuration write phase of the cli.
func TestWrite(t *testing.T) {
	cli := New().(*cliConfigurer)
	extendedCfg := Defaults()
	tf, _ := ioutil.TempFile("", "test_fritzctl.json.")
	defer tf.Close()
	defer os.Remove(tf.Name())
	extendedCfg.file = tf.Name()
	cli.ApplyDefaults(extendedCfg)
	err := cli.Write()
	assert.NoError(t, err)
}

// TestWriteAndRead test the configuration write with subsequent re-read.
func TestWriteAndRead(t *testing.T) {
	cli := New().(*cliConfigurer)
	extendedCfg := Defaults()
	tf, _ := ioutil.TempFile("", "test_fritzctl.json.")
	defer tf.Close()
	defer os.Remove(tf.Name())
	extendedCfg.file = tf.Name()
	cli.ApplyDefaults(extendedCfg)
	err := cli.Write()
	assert.NoError(t, err)
	re, err := config.New(tf.Name())
	assert.NoError(t, err)
	assert.NotNil(t, re)
	assert.Equal(t, *cli.userValues.fritzCfg.Net, *re.Net)
	assert.Equal(t, *cli.userValues.fritzCfg.Login, *re.Login)
	assert.Equal(t, *cli.userValues.fritzCfg.Pki, *re.Pki)
}

// TestWriteWithIOError test the write phase of the cli with error.
func TestWriteWithIOError(t *testing.T) {
	cli := New().(*cliConfigurer)
	extendedCfg := Defaults()
	extendedCfg.file = "/root/a/b/c/no/such/file/or/directory/cfg.json"
	cli.ApplyDefaults(extendedCfg)
	err := cli.Write()
	assert.Error(t, err)
}
<|endoftext|>


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

package service

import (
	"fmt"
	"io"
	"net"
	"time"

	"github.com/surgemq/message"
)

type netReader interface {
	io.Reader
	SetReadDeadline(t time.Time) error
}

type timeoutReader struct {
	d    time.Duration
	conn netReader
}

func (r timeoutReader) Read(b []byte) (int, error) {
	if err := r.conn.SetReadDeadline(time.Now().Add(r.d)); err != nil {
		return 0, err
	}
	return r.conn.Read(b)
}

// receiver() reads data from the network, and writes the data into the incoming buffer
func (this *service) receiver() {
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) receiver开始", this.cid())
	})
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic: %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping receiver", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting receiver", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		//Log.Debugc(func() string{ return fmt.Sprintf("server/handleConnection: Setting read deadline to %d", time.Second*time.Duration(this.keepAlive))})
		keepAlive := time.Second * time.Duration(this.keepAlive)
		r := timeoutReader{
			d:    keepAlive + (keepAlive / 2),
			conn: conn,
		}

		for {
			_, err := this.in.ReadFrom(r)
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sess is: %v", this.sess)})
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sessMgr is: %v", this.sessMgr)})

			/*if err != nil {
				Log.Infoc(func() string { return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err) })
				//         if err != io.EOF {
				//         }
				return
			}*/
			if err != nil {

				if err == io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err)
					})
					return
				} else {
					Log.Infoc(func() string {
						return fmt.Sprintf("(%s) info reading from connection: %v", this.cid(), err)
					})
				}

				continue
			} else {
				Log.Infoc(func() string {
					return fmt.Sprintf("(%s)向ringbuffer些数据成功！", this.cid())
				})
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket: %v", this.cid(), ErrInvalidConnectionType)})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) %v", this.cid(), ErrInvalidConnectionType)
		})
	}
}

// sender() writes data from the outgoing buffer to the network
func (this *service) sender() {
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) sender开始", this.cid())
	})
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic: %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping sender", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting sender", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		for {
			_, err := this.out.WriteTo(conn)

			if err != nil {
				if err != io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error writing data: %v", this.cid(), err)
					})
				}
				return
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket not supported", this.cid())})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Invalid connection type", this.cid())
		})
	}
}

// peekMessageSize() reads, but not commits, enough bytes to determine the size of
// the next message and returns the type and size.
func (this *service) peekMessageSize() (message.Message, int, error) {
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) peekMessageSize开始", this.cid())
	})
	var (
		b   *[]byte
		err error
		cnt int = 2
	)

	if this.in == nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize this.in is nil", this.cid())
		})
		err = ErrBufferNotReady
		return nil, 0, err
	}

	// Let's read enough bytes to get the message header (msg type, remaining length)
	//for {
	// If we have read 5 bytes and still not done, then there's a problem.
	//if cnt > 5 {
	//	return 0, 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
	//}

	// Peek cnt bytes from the input buffer.
	b, err = this.in.ReadWait(cnt)
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize this.in.ReadWait falure", this.cid())
		})
		return nil, 0, err
	}
	//fmt.Println("b=", b)
	//// If not enough bytes are returned, then continue until there's enough.
	//if len(b) < cnt {
	//	continue
	//}

	// If we got enough bytes, then check the last byte to see if the continuation
	// bit is set. If so, increment cnt and continue peeking
	/*if b[cnt-1] >= 0x80 {
		cnt++
	} else {
		break
	}*/
	//for {
	//	if cnt > 5 {
	//		return 0, 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
	//	}
	//
	//	if b[cnt-1] >= 0x80 {
	//		cnt++
	//	} else {
	//		break
	//	}
	//}
	//}

	// Get the remaining length of the message
	//remlen, m := binary.Uvarint(b[1:cnt])

	// Total message length is remlen + 1 (msg type) + m (remlen bytes)
	//total := int(remlen) + 1 + m

	mtype := message.MessageType(b[0] >> 4)

	//return mtype, total, err
	var msg message.Message
	msg, err = mtype.New()
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize mtype.New() falure", this.cid())
		})
		return nil, 0, err
	}
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) 开始创建对象(%s)", this.cid(), msg.Name())
	})
	_, err = msg.Decode(b)
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize msg.Decode falure", this.cid())
		})
		return nil, 0, err
	}
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) peekMessageSize结束(%s)", this.cid(), msg.Name())
	})
	return msg, len(b), err
}

// peekMessage() reads a message from the buffer, but the bytes are NOT committed.
// This means the buffer still thinks the bytes are not read yet.
/*func (this *service) peekMessage(mtype message.MessageType, total int) (message.Message, int, error) {
	var (
		b    []byte
		err  error
		i, n int
		msg  message.Message
	)

	if this.in == nil {
		return nil, 0, ErrBufferNotReady
	}

	// Peek until we get total bytes
	for i = 0; ; i++ {
		// Peek remlen bytes from the input buffer.
		b, err = this.in.ReadWait(total)
		if err != nil && err != ErrBufferInsufficientData {
			return nil, 0, err
		}

		// If not enough bytes are returned, then continue until there's enough.
		if len(b) >= total {
			break
		}
	}

	msg, err = mtype.New()
	if err != nil {
		return nil, 0, err
	}

	n, err = msg.Decode(b)
	return msg, n, err
}*/

// readMessage() reads and copies a message from the buffer. The buffer bytes are
// committed as a result of the read.
func (this *service) readMessage(mtype message.MessageType, total int) (message.Message, int, error) {
	var (
		b   []byte
		err error
		n   int
		msg message.Message
	)

	if this.in == nil {
		err = ErrBufferNotReady
		return nil, 0, err
	}

	if len(this.intmp) < total {
		this.intmp = make([]byte, total)
	}

	// Read until we get total bytes
	l := 0
	for l < total {
		n, err = this.in.Read(this.intmp[l:])
		l += n
		Log.Debugc(func() string {
			return fmt.Sprintf("read %d bytes, total %d", n, l)
		})
		if err != nil {
			return nil, 0, err
		}
	}

	b = this.intmp[:total]

	msg, err = mtype.New()
	if err != nil {
		return msg, 0, err
	}

	n, err = msg.Decode(b)
	return msg, n, err
}

// writeMessage() writes a message to the outgoing buffer
func (this *service) writeMessage(msg message.Message) (int, error) {
	var (
		l    int = msg.Len()
		m, n int
		err  error
		buf  *[]byte
		wrap bool
	)

	if this.out == nil {
		return 0, ErrBufferNotReady
	}

	// This is to serialize writes to the underlying buffer. Multiple goroutines could
	// potentially get here because of calling Publish() or Subscribe() or other
	// functions that will send messages. For example, if a message is received in
	// another connetion, and the message needs to be published to this client, then
	// the Publish() function is called, and at the same time, another client could
	// do exactly the same thing.
	//
	// Not an ideal fix though. If possible we should remove mutex and be lockfree.
	// Mainly because when there's a large number of goroutines that want to publish
	// to this client, then they will all block. However, this will do for now.
	//
	// FIXME: Try to find a better way than a mutex...if possible.
	this.wmu.Lock()
	defer this.wmu.Unlock()

	buf, wrap, err = this.out.WriteWait(l)
	if err != nil {
		return 0, err
	}

	if wrap {
		if len(this.outtmp) < l {
			this.outtmp = make([]byte, l)
		}

		n, err = msg.Encode(this.outtmp[0:])
		if err != nil {
			return 0, err
		}

		m, err = this.out.Write(this.outtmp[0:n])
		if err != nil {
			return m, err
		}
	} else {
		n, err = msg.Encode(buf[0:])
		if err != nil {
			return 0, err
		}

		m, err = this.out.WriteCommit(n)
		if err != nil {
			return 0, err
		}
	}

	this.outStat.increment(int64(m))

	return m, nil
}
<commit_msg>修改了buffer.go;sendrecv.go;process.go<commit_after>// Copyright (c) 2014 The SurgeMQ Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package service

import (
	"fmt"
	"io"
	"net"
	"time"

	"github.com/surgemq/message"
)

type netReader interface {
	io.Reader
	SetReadDeadline(t time.Time) error
}

type timeoutReader struct {
	d    time.Duration
	conn netReader
}

func (r timeoutReader) Read(b []byte) (int, error) {
	if err := r.conn.SetReadDeadline(time.Now().Add(r.d)); err != nil {
		return 0, err
	}
	return r.conn.Read(b)
}

// receiver() reads data from the network, and writes the data into the incoming buffer
func (this *service) receiver() {
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) receiver开始", this.cid())
	})
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic: %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping receiver", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting receiver", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		//Log.Debugc(func() string{ return fmt.Sprintf("server/handleConnection: Setting read deadline to %d", time.Second*time.Duration(this.keepAlive))})
		keepAlive := time.Second * time.Duration(this.keepAlive)
		r := timeoutReader{
			d:    keepAlive + (keepAlive / 2),
			conn: conn,
		}

		for {
			_, err := this.in.ReadFrom(r)
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sess is: %v", this.sess)})
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sessMgr is: %v", this.sessMgr)})

			/*if err != nil {
				Log.Infoc(func() string { return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err) })
				//         if err != io.EOF {
				//         }
				return
			}*/
			if err != nil {

				if err == io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err)
					})
					return
				} else {
					Log.Infoc(func() string {
						return fmt.Sprintf("(%s) info reading from connection: %v", this.cid(), err)
					})
				}

				continue
			} else {
				Log.Infoc(func() string {
					return fmt.Sprintf("(%s)向ringbuffer些数据成功！", this.cid())
				})
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket: %v", this.cid(), ErrInvalidConnectionType)})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) %v", this.cid(), ErrInvalidConnectionType)
		})
	}
}

// sender() writes data from the outgoing buffer to the network
func (this *service) sender() {
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) sender开始", this.cid())
	})
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic: %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping sender", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting sender", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		for {
			_, err := this.out.WriteTo(conn)

			if err != nil {
				if err != io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error writing data: %v", this.cid(), err)
					})
				}
				return
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket not supported", this.cid())})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Invalid connection type", this.cid())
		})
	}
}

// peekMessageSize() reads, but not commits, enough bytes to determine the size of
// the next message and returns the type and size.
func (this *service) peekMessageSize() (message.Message, int, error) {
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) peekMessageSize开始", this.cid())
	})
	var (
		b   *[]byte
		err error
		cnt int = 2
	)

	if this.in == nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize this.in is nil", this.cid())
		})
		err = ErrBufferNotReady
		return nil, 0, err
	}

	// Let's read enough bytes to get the message header (msg type, remaining length)
	//for {
	// If we have read 5 bytes and still not done, then there's a problem.
	//if cnt > 5 {
	//	return 0, 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
	//}

	// Peek cnt bytes from the input buffer.
	b, err = this.in.ReadWait(cnt)
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize this.in.ReadWait falure", this.cid())
		})
		return nil, 0, err
	}
	//fmt.Println("b=", b)
	//// If not enough bytes are returned, then continue until there's enough.
	//if len(b) < cnt {
	//	continue
	//}

	// If we got enough bytes, then check the last byte to see if the continuation
	// bit is set. If so, increment cnt and continue peeking
	/*if b[cnt-1] >= 0x80 {
		cnt++
	} else {
		break
	}*/
	//for {
	//	if cnt > 5 {
	//		return 0, 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
	//	}
	//
	//	if b[cnt-1] >= 0x80 {
	//		cnt++
	//	} else {
	//		break
	//	}
	//}
	//}

	// Get the remaining length of the message
	//remlen, m := binary.Uvarint(b[1:cnt])

	// Total message length is remlen + 1 (msg type) + m (remlen bytes)
	//total := int(remlen) + 1 + m

	mtype := message.MessageType(*b[0] >> 4)

	//return mtype, total, err
	var msg message.Message
	msg, err = mtype.New()
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize mtype.New() falure", this.cid())
		})
		return nil, 0, err
	}
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) 开始创建对象(%s)", this.cid(), msg.Name())
	})
	_, err = msg.Decode(*b)
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize msg.Decode falure", this.cid())
		})
		return nil, 0, err
	}
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) peekMessageSize结束(%s)", this.cid(), msg.Name())
	})
	return msg, len(*b), err
}

// peekMessage() reads a message from the buffer, but the bytes are NOT committed.
// This means the buffer still thinks the bytes are not read yet.
/*func (this *service) peekMessage(mtype message.MessageType, total int) (message.Message, int, error) {
	var (
		b    []byte
		err  error
		i, n int
		msg  message.Message
	)

	if this.in == nil {
		return nil, 0, ErrBufferNotReady
	}

	// Peek until we get total bytes
	for i = 0; ; i++ {
		// Peek remlen bytes from the input buffer.
		b, err = this.in.ReadWait(total)
		if err != nil && err != ErrBufferInsufficientData {
			return nil, 0, err
		}

		// If not enough bytes are returned, then continue until there's enough.
		if len(b) >= total {
			break
		}
	}

	msg, err = mtype.New()
	if err != nil {
		return nil, 0, err
	}

	n, err = msg.Decode(b)
	return msg, n, err
}*/

// readMessage() reads and copies a message from the buffer. The buffer bytes are
// committed as a result of the read.
func (this *service) readMessage(mtype message.MessageType, total int) (message.Message, int, error) {
	var (
		b   []byte
		err error
		n   int
		msg message.Message
	)

	if this.in == nil {
		err = ErrBufferNotReady
		return nil, 0, err
	}

	if len(this.intmp) < total {
		this.intmp = make([]byte, total)
	}

	// Read until we get total bytes
	l := 0
	for l < total {
		n, err = this.in.Read(this.intmp[l:])
		l += n
		Log.Debugc(func() string {
			return fmt.Sprintf("read %d bytes, total %d", n, l)
		})
		if err != nil {
			return nil, 0, err
		}
	}

	b = this.intmp[:total]

	msg, err = mtype.New()
	if err != nil {
		return msg, 0, err
	}

	n, err = msg.Decode(b)
	return msg, n, err
}

// writeMessage() writes a message to the outgoing buffer
func (this *service) writeMessage(msg message.Message) (int, error) {
	var (
		l    int = msg.Len()
		m, n int
		err  error
		buf  *[]byte
		wrap bool
	)

	if this.out == nil {
		return 0, ErrBufferNotReady
	}

	// This is to serialize writes to the underlying buffer. Multiple goroutines could
	// potentially get here because of calling Publish() or Subscribe() or other
	// functions that will send messages. For example, if a message is received in
	// another connetion, and the message needs to be published to this client, then
	// the Publish() function is called, and at the same time, another client could
	// do exactly the same thing.
	//
	// Not an ideal fix though. If possible we should remove mutex and be lockfree.
	// Mainly because when there's a large number of goroutines that want to publish
	// to this client, then they will all block. However, this will do for now.
	//
	// FIXME: Try to find a better way than a mutex...if possible.
	this.wmu.Lock()
	defer this.wmu.Unlock()

	buf, wrap, err = this.out.WriteWait(l)
	if err != nil {
		return 0, err
	}

	if wrap {
		if len(this.outtmp) < l {
			this.outtmp = make([]byte, l)
		}

		n, err = msg.Encode(this.outtmp[0:])
		if err != nil {
			return 0, err
		}

		m, err = this.out.Write(this.outtmp[0:n])
		if err != nil {
			return m, err
		}
	} else {
		n, err = msg.Encode(*buf[0:])
		if err != nil {
			return 0, err
		}

		m, err = this.out.WriteCommit(n)
		if err != nil {
			return 0, err
		}
	}

	this.outStat.increment(int64(m))

	return m, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"os/exec"

	"github.com/gorilla/websocket"
	"github.com/pborman/uuid"

	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/logger"
)

// rsyncCopy copies a directory using rsync (with the --devices option).
func rsyncLocalCopy(source string, dest string, bwlimit string) (string, error) {
	err := os.MkdirAll(dest, 0755)
	if err != nil {
		return "", err
	}

	rsyncVerbosity := "-q"
	if debug {
		rsyncVerbosity = "-vi"
	}

	if bwlimit == "" {
		bwlimit = "0"
	}

	return shared.RunCommand("rsync",
		"-a",
		"-HAX",
		"--sparse",
		"--devices",
		"--delete",
		"--checksum",
		"--numeric-ids",
		"--bwlimit", bwlimit,
		rsyncVerbosity,
		shared.AddSlash(source),
		dest)
}

func rsyncSendSetup(name string, path string, bwlimit string) (*exec.Cmd, net.Conn, io.ReadCloser, error) {
	/*
	 * The way rsync works, it invokes a subprocess that does the actual
	 * talking (given to it by a -E argument). Since there isn't an easy
	 * way for us to capture this process' stdin/stdout, we just use netcat
	 * and write to/from a unix socket.
	 *
	 * In principle we don't need this socket. It seems to me that some
	 * clever invocation of rsync --server --sender and usage of that
	 * process' stdin/stdout could work around the need for this socket,
	 * but I couldn't get it to work. Another option would be to look at
	 * the spawned process' first child and read/write from its
	 * stdin/stdout, but that also seemed messy. In any case, this seems to
	 * work just fine.
	 */
	auds := fmt.Sprintf("@lxd/%s", uuid.NewRandom().String())
	// We simply copy a part of the uuid if it's longer than the allowed
	// maximum. That should be safe enough for our purposes.
	if len(auds) > shared.ABSTRACT_UNIX_SOCK_LEN-1 {
		auds = auds[:shared.ABSTRACT_UNIX_SOCK_LEN-1]
	}
	l, err := net.Listen("unix", auds)
	if err != nil {
		return nil, nil, nil, err
	}

	/*
	 * Here, the path /tmp/foo is ignored. Since we specify localhost,
	 * rsync thinks we are syncing to a remote host (in this case, the
	 * other end of the lxd websocket), and so the path specified on the
	 * --server instance of rsync takes precedence.
	 *
	 * Additionally, we use sh -c instead of just calling nc directly
	 * because rsync passes a whole bunch of arguments to the wrapper
	 * command (i.e. the command to run on --server). However, we're
	 * hardcoding that at the other end, so we can just ignore it.
	 */
	rsyncCmd := fmt.Sprintf("sh -c \"%s netcat %s %s\"", execPath, auds, name)
	if bwlimit == "" {
		bwlimit = "0"
	}

	cmd := exec.Command("rsync",
		"-arvP",
		"--devices",
		"--numeric-ids",
		"--partial",
		"--sparse",
		path,
		"localhost:/tmp/foo",
		"-e",
		rsyncCmd,
		"--bwlimit",
		bwlimit)

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return nil, nil, nil, err
	}

	if err := cmd.Start(); err != nil {
		return nil, nil, nil, err
	}

	conn, err := l.Accept()
	if err != nil {
		cmd.Process.Kill()
		cmd.Wait()
		return nil, nil, nil, err
	}
	l.Close()

	return cmd, conn, stderr, nil
}

// RsyncSend sets up the sending half of an rsync, to recursively send the
// directory pointed to by path over the websocket.
func RsyncSend(name string, path string, conn *websocket.Conn, readWrapper func(io.ReadCloser) io.ReadCloser, bwlimit string) error {
	cmd, dataSocket, stderr, err := rsyncSendSetup(name, path, bwlimit)
	if err != nil {
		return err
	}

	if dataSocket != nil {
		defer dataSocket.Close()
	}

	readPipe := io.ReadCloser(dataSocket)
	if readWrapper != nil {
		readPipe = readWrapper(dataSocket)
	}

	readDone, writeDone := shared.WebsocketMirror(conn, dataSocket, readPipe, nil, nil)

	output, err := ioutil.ReadAll(stderr)
	if err != nil {
		cmd.Process.Kill()
		cmd.Wait()
		return err
	}

	err = cmd.Wait()
	if err != nil {
		logger.Errorf("Rsync send failed: %s: %s: %s", path, err, string(output))
	}

	<-readDone
	<-writeDone

	return err
}

// RsyncRecv sets up the receiving half of the websocket to rsync (the other
// half set up by RsyncSend), putting the contents in the directory specified
// by path.
func RsyncRecv(path string, conn *websocket.Conn, writeWrapper func(io.WriteCloser) io.WriteCloser) error {
	cmd := exec.Command("rsync",
		"--server",
		"-vlogDtpre.iLsfx",
		"--numeric-ids",
		"--devices",
		"--partial",
		"--sparse",
		".",
		path)

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return err
	}

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return err
	}

	if err := cmd.Start(); err != nil {
		return err
	}

	writePipe := io.WriteCloser(stdin)
	if writeWrapper != nil {
		writePipe = writeWrapper(stdin)
	}

	readDone, writeDone := shared.WebsocketMirror(conn, writePipe, stdout, nil, nil)
	output, err := ioutil.ReadAll(stderr)
	if err != nil {
		cmd.Process.Kill()
		cmd.Wait()
		return err
	}

	err = cmd.Wait()
	if err != nil {
		logger.Errorf("Rsync receive failed: %s: %s: %s", path, err, string(output))
	}

	<-readDone
	<-writeDone

	return err
}
<commit_msg>rsync: ignore vanished file warnings<commit_after>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"syscall"

	"github.com/gorilla/websocket"
	"github.com/pborman/uuid"

	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/logger"
)

// rsyncCopy copies a directory using rsync (with the --devices option).
func rsyncLocalCopy(source string, dest string, bwlimit string) (string, error) {
	err := os.MkdirAll(dest, 0755)
	if err != nil {
		return "", err
	}

	rsyncVerbosity := "-q"
	if debug {
		rsyncVerbosity = "-vi"
	}

	if bwlimit == "" {
		bwlimit = "0"
	}

	msg, err := shared.RunCommand("rsync",
		"-a",
		"-HAX",
		"--sparse",
		"--devices",
		"--delete",
		"--checksum",
		"--numeric-ids",
		"--bwlimit", bwlimit,
		rsyncVerbosity,
		shared.AddSlash(source),
		dest)
	if err != nil {
		runError, ok := err.(shared.RunError)
		if ok {
			exitError, ok := runError.Err.(*exec.ExitError)
			if ok {
				waitStatus := exitError.Sys().(syscall.WaitStatus)
				if waitStatus.ExitStatus() == 24 {
					return msg, nil
				}
			}
		}
		return msg, err
	}

	return msg, nil
}

func rsyncSendSetup(name string, path string, bwlimit string) (*exec.Cmd, net.Conn, io.ReadCloser, error) {
	/*
	 * The way rsync works, it invokes a subprocess that does the actual
	 * talking (given to it by a -E argument). Since there isn't an easy
	 * way for us to capture this process' stdin/stdout, we just use netcat
	 * and write to/from a unix socket.
	 *
	 * In principle we don't need this socket. It seems to me that some
	 * clever invocation of rsync --server --sender and usage of that
	 * process' stdin/stdout could work around the need for this socket,
	 * but I couldn't get it to work. Another option would be to look at
	 * the spawned process' first child and read/write from its
	 * stdin/stdout, but that also seemed messy. In any case, this seems to
	 * work just fine.
	 */
	auds := fmt.Sprintf("@lxd/%s", uuid.NewRandom().String())
	// We simply copy a part of the uuid if it's longer than the allowed
	// maximum. That should be safe enough for our purposes.
	if len(auds) > shared.ABSTRACT_UNIX_SOCK_LEN-1 {
		auds = auds[:shared.ABSTRACT_UNIX_SOCK_LEN-1]
	}
	l, err := net.Listen("unix", auds)
	if err != nil {
		return nil, nil, nil, err
	}

	/*
	 * Here, the path /tmp/foo is ignored. Since we specify localhost,
	 * rsync thinks we are syncing to a remote host (in this case, the
	 * other end of the lxd websocket), and so the path specified on the
	 * --server instance of rsync takes precedence.
	 *
	 * Additionally, we use sh -c instead of just calling nc directly
	 * because rsync passes a whole bunch of arguments to the wrapper
	 * command (i.e. the command to run on --server). However, we're
	 * hardcoding that at the other end, so we can just ignore it.
	 */
	rsyncCmd := fmt.Sprintf("sh -c \"%s netcat %s %s\"", execPath, auds, name)
	if bwlimit == "" {
		bwlimit = "0"
	}

	cmd := exec.Command("rsync",
		"-arvP",
		"--devices",
		"--numeric-ids",
		"--partial",
		"--sparse",
		path,
		"localhost:/tmp/foo",
		"-e",
		rsyncCmd,
		"--bwlimit",
		bwlimit)

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return nil, nil, nil, err
	}

	if err := cmd.Start(); err != nil {
		return nil, nil, nil, err
	}

	conn, err := l.Accept()
	if err != nil {
		cmd.Process.Kill()
		cmd.Wait()
		return nil, nil, nil, err
	}
	l.Close()

	return cmd, conn, stderr, nil
}

// RsyncSend sets up the sending half of an rsync, to recursively send the
// directory pointed to by path over the websocket.
func RsyncSend(name string, path string, conn *websocket.Conn, readWrapper func(io.ReadCloser) io.ReadCloser, bwlimit string) error {
	cmd, dataSocket, stderr, err := rsyncSendSetup(name, path, bwlimit)
	if err != nil {
		return err
	}

	if dataSocket != nil {
		defer dataSocket.Close()
	}

	readPipe := io.ReadCloser(dataSocket)
	if readWrapper != nil {
		readPipe = readWrapper(dataSocket)
	}

	readDone, writeDone := shared.WebsocketMirror(conn, dataSocket, readPipe, nil, nil)

	output, err := ioutil.ReadAll(stderr)
	if err != nil {
		cmd.Process.Kill()
		cmd.Wait()
		return err
	}

	err = cmd.Wait()
	if err != nil {
		logger.Errorf("Rsync send failed: %s: %s: %s", path, err, string(output))
	}

	<-readDone
	<-writeDone

	return err
}

// RsyncRecv sets up the receiving half of the websocket to rsync (the other
// half set up by RsyncSend), putting the contents in the directory specified
// by path.
func RsyncRecv(path string, conn *websocket.Conn, writeWrapper func(io.WriteCloser) io.WriteCloser) error {
	cmd := exec.Command("rsync",
		"--server",
		"-vlogDtpre.iLsfx",
		"--numeric-ids",
		"--devices",
		"--partial",
		"--sparse",
		".",
		path)

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return err
	}

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return err
	}

	if err := cmd.Start(); err != nil {
		return err
	}

	writePipe := io.WriteCloser(stdin)
	if writeWrapper != nil {
		writePipe = writeWrapper(stdin)
	}

	readDone, writeDone := shared.WebsocketMirror(conn, writePipe, stdout, nil, nil)
	output, err := ioutil.ReadAll(stderr)
	if err != nil {
		cmd.Process.Kill()
		cmd.Wait()
		return err
	}

	err = cmd.Wait()
	if err != nil {
		logger.Errorf("Rsync receive failed: %s: %s: %s", path, err, string(output))
	}

	<-readDone
	<-writeDone

	return err
}
<|endoftext|>


<commit_before>// Copyright 2017 gf Author(https://github.com/gogf/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

// Package gfile provides easy-to-use operations for file system.
package gfile

import (
	"bytes"
	"errors"
	"fmt"
	"github.com/gogf/gf/g/container/gtype"
	"github.com/gogf/gf/g/text/gregex"
	"github.com/gogf/gf/g/text/gstr"
	"github.com/gogf/gf/g/util/gconv"
	"io"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
	"time"
)

const (
    // Separator for file system.
    Separator     = string(filepath.Separator)
    // Default perm for file opening.
    gDEFAULT_PERM = 0666
)

var (
    // The absolute file path for main package.
    // It can be only checked and set once.
    mainPkgPath = gtype.NewString()
)

// Mkdir creates directories recursively with given <path>.
// The parameter <path> is suggested to be absolute path.
func Mkdir(path string) error {
    err  := os.MkdirAll(path, os.ModePerm)
    if err != nil {
        return err
    }
    return nil
}

// Create creates file with given <path> recursively.
// The parameter <path> is suggested to be absolute path.
func Create(path string) (*os.File, error) {
    dir := Dir(path)
    if !Exists(dir) {
        Mkdir(dir)
    }
    return os.Create(path)
}

// Open opens file/directory readonly.
func Open(path string) (*os.File, error) {
    return os.Open(path)
}

// OpenFile opens file/directory with given <flag> and <perm>.
func OpenFile(path string, flag int, perm os.FileMode) (*os.File, error) {
    return os.OpenFile(path, flag, perm)
}

// OpenWithFlag opens file/directory with default perm and given <flag>.
func OpenWithFlag(path string, flag int) (*os.File, error) {
    f, err  := os.OpenFile(path, flag, gDEFAULT_PERM)
    if err != nil {
        return nil, err
    }
    return f, nil
}

// OpenWithFlagPerm opens file/directory with given <flag> and <perm>.
func OpenWithFlagPerm(path string, flag int, perm int) (*os.File, error) {
    f, err  := os.OpenFile(path, flag, os.FileMode(perm))
    if err != nil {
        return nil, err
    }
    return f, nil
}

// Exists checks whether given <path> exist.
func Exists(path string) bool {
    if _, err := os.Stat(path); !os.IsNotExist(err) {
        return true
    }
    return false
}

// IsDir checks whether given <path> a directory.
func IsDir(path string) bool {
    s, err := os.Stat(path)
    if err != nil {
        return false
    }
    return s.IsDir()
}

// Pwd returns absolute path of current working directory.
func Pwd() string {
    path, _ := os.Getwd()
    return path
}

// IsFile checks whether given <path> a file, which means it's not a directory.
func IsFile(path string) bool {
    s, err := os.Stat(path)
    if err != nil {
        return false
    }
    return !s.IsDir()
}

// Alias of Stat.
// See Stat.
func Info(path string) (os.FileInfo, error) {
    return Stat(path)
}

// Stat returns a FileInfo describing the named file.
// If there is an error, it will be of type *PathError.
func Stat(path string) (os.FileInfo, error) {
    return os.Stat(path)
}

// Move renames (moves) <src> to <dst> path.
func Move(src string, dst string) error {
    return os.Rename(src, dst)
}

// Alias of Move.
// See Move.
func Rename(src string, dst string) error {
    return Move(src, dst)
}

// Copy file from <src> to <dst>.
//
// @TODO directory copy support.
func Copy(src string, dst string) error {
    srcFile, err := Open(src)
    if err != nil {
        return err
    }
    defer srcFile.Close()
    dstFile, err := Create(dst)
    if err != nil {
        return err
    }
    defer dstFile.Close()
    _, err = io.Copy(dstFile, srcFile)
    if err != nil {
        return err
    }
    err = dstFile.Sync()
    if err != nil {
        return err
    }
    return nil
}

// DirNames returns sub-file names of given directory <path>.
func DirNames(path string) ([]string, error) {
    f, err := os.Open(path)
    if err != nil {
        return nil, err
    }
    list, err := f.Readdirnames(-1)
    f.Close()
    if err != nil {
        return nil, err
    }
    return list, nil
}

// Glob returns the names of all files matching pattern or nil
// if there is no matching file. The syntax of patterns is the same
// as in Match. The pattern may describe hierarchical names such as
// /usr/*/bin/ed (assuming the Separator is '/').
//
// Glob ignores file system errors such as I/O errors reading directories.
// The only possible returned error is ErrBadPattern, when pattern
// is malformed.
func Glob(pattern string, onlyNames...bool) ([]string, error) {
    if list, err := filepath.Glob(pattern); err == nil {
        if len(onlyNames) > 0 && onlyNames[0] && len(list) > 0 {
            array := make([]string, len(list))
            for k, v := range list {
                array[k] = Basename(v)
            }
            return array, nil
        }
        return list, nil
    } else {
        return nil, err
    }
}

// Remove deletes all file/directory with <path> parameter.
// If parameter <path> is directory, it deletes it recursively.
func Remove(path string) error {
    return os.RemoveAll(path)
}

// IsReadable checks whether given <path> is readable.
func IsReadable(path string) bool {
    result    := true
    file, err := os.OpenFile(path, os.O_RDONLY, gDEFAULT_PERM)
    if err != nil {
        result = false
    }
    file.Close()
    return result
}

// IsWritable checks whether given <path> is writable.
//
// @TODO improve performance; use golang.org/x/sys to cross-plat-form
func IsWritable(path string) bool {
    result := true
    if IsDir(path) {
        // If it's a directory, create a temporary file to test whether it's writable.
        tmpFile := strings.TrimRight(path, Separator) + Separator + gconv.String(time.Now().UnixNano())
        if f, err := Create(tmpFile); err != nil || !Exists(tmpFile){
            result = false
        } else {
            f.Close()
            Remove(tmpFile)
        }
    } else {
        // 如果是文件，那么判断文件是否可打开
        file, err := os.OpenFile(path, os.O_WRONLY, gDEFAULT_PERM)
        if err != nil {
            result = false
        }
        file.Close()
    }
    return result
}

// See os.Chmod.
func Chmod(path string, mode os.FileMode) error {
    return os.Chmod(path, mode)
}

// ScanDir returns all sub-files with absolute paths of given <path>,
// It scans directory recursively if given parameter <recursive> is true.
func ScanDir(path string, pattern string, recursive ... bool) ([]string, error) {
    list, err := doScanDir(path, pattern, recursive...)
    if err != nil {
        return nil, err
    }
    if len(list) > 0 {
        sort.Strings(list)
    }
    return list, nil
}

// doScanDir is an internal method which scans directory
// and returns the absolute path list of files that are not sorted.
//
// The pattern parameter <pattern> supports multiple file name patterns,
// using the ',' symbol to separate multiple patterns.
//
// It scans directory recursively if given parameter <recursive> is true.
func doScanDir(path string, pattern string, recursive ... bool) ([]string, error) {
    list := ([]string)(nil)
    file, err := os.Open(path)
    if err != nil {
        return nil, err
    }
    defer file.Close()
    names, err := file.Readdirnames(-1)
    if err != nil {
        return nil, err
    }
    for _, name := range names {
        path := fmt.Sprintf("%s%s%s", path, Separator, name)
        if IsDir(path) && len(recursive) > 0 && recursive[0] {
            array, _ := doScanDir(path, pattern, true)
            if len(array) > 0 {
                list = append(list, array...)
            }
        }
        // If it meets pattern, then add it to the result list.
        for _, p := range strings.Split(pattern, ",") {
            if match, err := filepath.Match(strings.TrimSpace(p), name); err == nil && match {
                list = append(list, path)
            }
        }
    }
    return list, nil
}

// RealPath converts the given <path> to its absolute path
// and checks if the file path exists.
// If the file does not exist, return an empty string.
func RealPath(path string) string {
    p, err := filepath.Abs(path)
    if err != nil {
        return ""
    }
    if !Exists(p) {
        return ""
    }
    return p
}

// SelfPath returns absolute file path of current running process(binary).
func SelfPath() string {
    p, _ := filepath.Abs(os.Args[0])
    return p
}

// SelfDir returns absolute directory path of current running process(binary).
func SelfDir() string {
    return filepath.Dir(SelfPath())
}

// Basename returns the last element of path.
// Trailing path separators are removed before extracting the last element.
// If the path is empty, Base returns ".".
// If the path consists entirely of separators, Base returns a single separator.
func Basename(path string) string {
    return filepath.Base(path)
}

// Dir returns all but the last element of path, typically the path's directory.
// After dropping the final element, Dir calls Clean on the path and trailing
// slashes are removed.
// If the path is empty, Dir returns ".".
// If the path consists entirely of separators, Dir returns a single separator.
// The returned path does not end in a separator unless it is the root directory.
func Dir(path string) string {
    return filepath.Dir(path)
}

// Ext returns the file name extension used by path.
// The extension is the suffix beginning at the final dot
// in the final element of path; it is empty if there is
// no dot.
//
// Note: the result contains symbol '.'.
func Ext(path string) string {
    return filepath.Ext(path)
}

// Home returns absolute path of current user's home directory.
func Home() (string, error) {
    u, err := user.Current()
    if nil == err {
        return u.HomeDir, nil
    }
    if "windows" == runtime.GOOS {
        return homeWindows()
    }
    return homeUnix()
}

func homeUnix() (string, error) {
    if home := os.Getenv("HOME"); home != "" {
        return home, nil
    }
    var stdout bytes.Buffer
    cmd := exec.Command("sh", "-c", "eval echo ~$USER")
    cmd.Stdout = &stdout
    if err := cmd.Run(); err != nil {
        return "", err
    }

    result := strings.TrimSpace(stdout.String())
    if result == "" {
        return "", errors.New("blank output when reading home directory")
    }

    return result, nil
}

func homeWindows() (string, error) {
    drive := os.Getenv("HOMEDRIVE")
    path  := os.Getenv("HOMEPATH")
    home  := drive + path
    if drive == "" || path == "" {
        home = os.Getenv("USERPROFILE")
    }
    if home == "" {
        return "", errors.New("HOMEDRIVE, HOMEPATH, and USERPROFILE are blank")
    }

    return home, nil
}

// MainPkgPath returns absolute file path of package main,
// which contains the entrance function main.
//
// It's only available in develop environment.
//
// Note1: Only valid for source development environments,
// IE only valid for systems that generate this executable.
// Note2: When the method is called for the first time, if it is in an asynchronous goroutine,
// the method may not get the main package path.
func MainPkgPath() string {
    path := mainPkgPath.Val()
    if path != "" {
    	if path == "-" {
    		return ""
	    }
        return path
    }
    for i := 1; i < 10000; i++ {
        if _, file, _, ok := runtime.Caller(i); ok {
        	// <file> is separated by '/'
        	if gstr.Contains(file, "/gf/g/") {
        		continue
	        }
        	if Ext(file) != ".go" {
        		continue
	        }
        	// separator of <file> '/' will be converted to Separator.
        	path = Dir(file)
        	for path[len(path) - 1] != os.PathSeparator {
        		files, _ := ScanDir(path, "*.go")
        		for _, v := range files {
			        if gregex.IsMatchString(`package\s+main`, GetContents(v)) {
				        mainPkgPath.Set(path)
				        return path
			        }
		        }
        		path = Dir(path)
	        }

        } else {
            break
        }
    }
    // If it fails finding the path, then mark it as "-",
    // which means it will never do this search again.
	mainPkgPath.Set("-")
    return ""
}

// See os.TempDir().
func TempDir() string {
    return os.TempDir()
}
<commit_msg>fix issue in gfile.MainPkgPath<commit_after>// Copyright 2017 gf Author(https://github.com/gogf/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

// Package gfile provides easy-to-use operations for file system.
package gfile

import (
	"bytes"
	"errors"
	"fmt"
	"github.com/gogf/gf/g/container/gtype"
	"github.com/gogf/gf/g/text/gregex"
	"github.com/gogf/gf/g/text/gstr"
	"github.com/gogf/gf/g/util/gconv"
	"io"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
	"time"
)

const (
    // Separator for file system.
    Separator     = string(filepath.Separator)
    // Default perm for file opening.
    gDEFAULT_PERM = 0666
)

var (
    // The absolute file path for main package.
    // It can be only checked and set once.
    mainPkgPath = gtype.NewString()
)

// Mkdir creates directories recursively with given <path>.
// The parameter <path> is suggested to be absolute path.
func Mkdir(path string) error {
    err  := os.MkdirAll(path, os.ModePerm)
    if err != nil {
        return err
    }
    return nil
}

// Create creates file with given <path> recursively.
// The parameter <path> is suggested to be absolute path.
func Create(path string) (*os.File, error) {
    dir := Dir(path)
    if !Exists(dir) {
        Mkdir(dir)
    }
    return os.Create(path)
}

// Open opens file/directory readonly.
func Open(path string) (*os.File, error) {
    return os.Open(path)
}

// OpenFile opens file/directory with given <flag> and <perm>.
func OpenFile(path string, flag int, perm os.FileMode) (*os.File, error) {
    return os.OpenFile(path, flag, perm)
}

// OpenWithFlag opens file/directory with default perm and given <flag>.
func OpenWithFlag(path string, flag int) (*os.File, error) {
    f, err  := os.OpenFile(path, flag, gDEFAULT_PERM)
    if err != nil {
        return nil, err
    }
    return f, nil
}

// OpenWithFlagPerm opens file/directory with given <flag> and <perm>.
func OpenWithFlagPerm(path string, flag int, perm int) (*os.File, error) {
    f, err  := os.OpenFile(path, flag, os.FileMode(perm))
    if err != nil {
        return nil, err
    }
    return f, nil
}

// Exists checks whether given <path> exist.
func Exists(path string) bool {
    if _, err := os.Stat(path); !os.IsNotExist(err) {
        return true
    }
    return false
}

// IsDir checks whether given <path> a directory.
func IsDir(path string) bool {
    s, err := os.Stat(path)
    if err != nil {
        return false
    }
    return s.IsDir()
}

// Pwd returns absolute path of current working directory.
func Pwd() string {
    path, _ := os.Getwd()
    return path
}

// IsFile checks whether given <path> a file, which means it's not a directory.
func IsFile(path string) bool {
    s, err := os.Stat(path)
    if err != nil {
        return false
    }
    return !s.IsDir()
}

// Alias of Stat.
// See Stat.
func Info(path string) (os.FileInfo, error) {
    return Stat(path)
}

// Stat returns a FileInfo describing the named file.
// If there is an error, it will be of type *PathError.
func Stat(path string) (os.FileInfo, error) {
    return os.Stat(path)
}

// Move renames (moves) <src> to <dst> path.
func Move(src string, dst string) error {
    return os.Rename(src, dst)
}

// Alias of Move.
// See Move.
func Rename(src string, dst string) error {
    return Move(src, dst)
}

// Copy file from <src> to <dst>.
//
// @TODO directory copy support.
func Copy(src string, dst string) error {
    srcFile, err := Open(src)
    if err != nil {
        return err
    }
    defer srcFile.Close()
    dstFile, err := Create(dst)
    if err != nil {
        return err
    }
    defer dstFile.Close()
    _, err = io.Copy(dstFile, srcFile)
    if err != nil {
        return err
    }
    err = dstFile.Sync()
    if err != nil {
        return err
    }
    return nil
}

// DirNames returns sub-file names of given directory <path>.
func DirNames(path string) ([]string, error) {
    f, err := os.Open(path)
    if err != nil {
        return nil, err
    }
    list, err := f.Readdirnames(-1)
    f.Close()
    if err != nil {
        return nil, err
    }
    return list, nil
}

// Glob returns the names of all files matching pattern or nil
// if there is no matching file. The syntax of patterns is the same
// as in Match. The pattern may describe hierarchical names such as
// /usr/*/bin/ed (assuming the Separator is '/').
//
// Glob ignores file system errors such as I/O errors reading directories.
// The only possible returned error is ErrBadPattern, when pattern
// is malformed.
func Glob(pattern string, onlyNames...bool) ([]string, error) {
    if list, err := filepath.Glob(pattern); err == nil {
        if len(onlyNames) > 0 && onlyNames[0] && len(list) > 0 {
            array := make([]string, len(list))
            for k, v := range list {
                array[k] = Basename(v)
            }
            return array, nil
        }
        return list, nil
    } else {
        return nil, err
    }
}

// Remove deletes all file/directory with <path> parameter.
// If parameter <path> is directory, it deletes it recursively.
func Remove(path string) error {
    return os.RemoveAll(path)
}

// IsReadable checks whether given <path> is readable.
func IsReadable(path string) bool {
    result    := true
    file, err := os.OpenFile(path, os.O_RDONLY, gDEFAULT_PERM)
    if err != nil {
        result = false
    }
    file.Close()
    return result
}

// IsWritable checks whether given <path> is writable.
//
// @TODO improve performance; use golang.org/x/sys to cross-plat-form
func IsWritable(path string) bool {
    result := true
    if IsDir(path) {
        // If it's a directory, create a temporary file to test whether it's writable.
        tmpFile := strings.TrimRight(path, Separator) + Separator + gconv.String(time.Now().UnixNano())
        if f, err := Create(tmpFile); err != nil || !Exists(tmpFile){
            result = false
        } else {
            f.Close()
            Remove(tmpFile)
        }
    } else {
        // 如果是文件，那么判断文件是否可打开
        file, err := os.OpenFile(path, os.O_WRONLY, gDEFAULT_PERM)
        if err != nil {
            result = false
        }
        file.Close()
    }
    return result
}

// See os.Chmod.
func Chmod(path string, mode os.FileMode) error {
    return os.Chmod(path, mode)
}

// ScanDir returns all sub-files with absolute paths of given <path>,
// It scans directory recursively if given parameter <recursive> is true.
func ScanDir(path string, pattern string, recursive ... bool) ([]string, error) {
    list, err := doScanDir(path, pattern, recursive...)
    if err != nil {
        return nil, err
    }
    if len(list) > 0 {
        sort.Strings(list)
    }
    return list, nil
}

// doScanDir is an internal method which scans directory
// and returns the absolute path list of files that are not sorted.
//
// The pattern parameter <pattern> supports multiple file name patterns,
// using the ',' symbol to separate multiple patterns.
//
// It scans directory recursively if given parameter <recursive> is true.
func doScanDir(path string, pattern string, recursive ... bool) ([]string, error) {
    list := ([]string)(nil)
    file, err := os.Open(path)
    if err != nil {
        return nil, err
    }
    defer file.Close()
    names, err := file.Readdirnames(-1)
    if err != nil {
        return nil, err
    }
    for _, name := range names {
        path := fmt.Sprintf("%s%s%s", path, Separator, name)
        if IsDir(path) && len(recursive) > 0 && recursive[0] {
            array, _ := doScanDir(path, pattern, true)
            if len(array) > 0 {
                list = append(list, array...)
            }
        }
        // If it meets pattern, then add it to the result list.
        for _, p := range strings.Split(pattern, ",") {
            if match, err := filepath.Match(strings.TrimSpace(p), name); err == nil && match {
                list = append(list, path)
            }
        }
    }
    return list, nil
}

// RealPath converts the given <path> to its absolute path
// and checks if the file path exists.
// If the file does not exist, return an empty string.
func RealPath(path string) string {
    p, err := filepath.Abs(path)
    if err != nil {
        return ""
    }
    if !Exists(p) {
        return ""
    }
    return p
}

// SelfPath returns absolute file path of current running process(binary).
func SelfPath() string {
    p, _ := filepath.Abs(os.Args[0])
    return p
}

// SelfDir returns absolute directory path of current running process(binary).
func SelfDir() string {
    return filepath.Dir(SelfPath())
}

// Basename returns the last element of path.
// Trailing path separators are removed before extracting the last element.
// If the path is empty, Base returns ".".
// If the path consists entirely of separators, Base returns a single separator.
func Basename(path string) string {
    return filepath.Base(path)
}

// Dir returns all but the last element of path, typically the path's directory.
// After dropping the final element, Dir calls Clean on the path and trailing
// slashes are removed.
// If the path is empty, Dir returns ".".
// If the path consists entirely of separators, Dir returns a single separator.
// The returned path does not end in a separator unless it is the root directory.
func Dir(path string) string {
    return filepath.Dir(path)
}

// Ext returns the file name extension used by path.
// The extension is the suffix beginning at the final dot
// in the final element of path; it is empty if there is
// no dot.
//
// Note: the result contains symbol '.'.
func Ext(path string) string {
    return filepath.Ext(path)
}

// Home returns absolute path of current user's home directory.
func Home() (string, error) {
    u, err := user.Current()
    if nil == err {
        return u.HomeDir, nil
    }
    if "windows" == runtime.GOOS {
        return homeWindows()
    }
    return homeUnix()
}

func homeUnix() (string, error) {
    if home := os.Getenv("HOME"); home != "" {
        return home, nil
    }
    var stdout bytes.Buffer
    cmd := exec.Command("sh", "-c", "eval echo ~$USER")
    cmd.Stdout = &stdout
    if err := cmd.Run(); err != nil {
        return "", err
    }

    result := strings.TrimSpace(stdout.String())
    if result == "" {
        return "", errors.New("blank output when reading home directory")
    }

    return result, nil
}

func homeWindows() (string, error) {
    drive := os.Getenv("HOMEDRIVE")
    path  := os.Getenv("HOMEPATH")
    home  := drive + path
    if drive == "" || path == "" {
        home = os.Getenv("USERPROFILE")
    }
    if home == "" {
        return "", errors.New("HOMEDRIVE, HOMEPATH, and USERPROFILE are blank")
    }

    return home, nil
}

// MainPkgPath returns absolute file path of package main,
// which contains the entrance function main.
//
// It's only available in develop environment.
//
// Note1: Only valid for source development environments,
// IE only valid for systems that generate this executable.
// Note2: When the method is called for the first time, if it is in an asynchronous goroutine,
// the method may not get the main package path.
func MainPkgPath() string {
    path := mainPkgPath.Val()
    if path != "" {
    	if path == "-" {
    		return ""
	    }
        return path
    }
    for i := 1; i < 10000; i++ {
        if _, file, _, ok := runtime.Caller(i); ok {
        	// <file> is separated by '/'
        	if gstr.Contains(file, "/gf/g/") {
        		continue
	        }
        	if Ext(file) != ".go" {
        		continue
	        }
        	// separator of <file> '/' will be converted to Separator.
        	for path = Dir(file); len(path) > 1 && Exists(path) && path[len(path) - 1] != os.PathSeparator; {
        		files, _ := ScanDir(path, "*.go")
        		for _, v := range files {
			        if gregex.IsMatchString(`package\s+main`, GetContents(v)) {
				        mainPkgPath.Set(path)
				        return path
			        }
		        }
        		path = Dir(path)
	        }

        } else {
            break
        }
    }
    // If it fails finding the path, then mark it as "-",
    // which means it will never do this search again.
	mainPkgPath.Set("-")
    return ""
}

// See os.TempDir().
func TempDir() string {
    return os.TempDir()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"os/signal"
	"syscall"

	"github.com/ReformedDevs/anonbot/db"
	"github.com/ReformedDevs/anonbot/server"
	"github.com/ReformedDevs/anonbot/tweeter"
	"github.com/howeyc/gopass"
	"github.com/urfave/cli"
)

func main() {
	app := cli.NewApp()
	app.Name = "anonbot"
	app.Usage = "Twitter anon account bot"
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "db-args",
			Value:  "dbname=postgres user=postgres",
			EnvVar: "DB_ARGS",
			Usage:  "database arguments",
		},
		cli.StringFlag{
			Name:   "db-driver",
			Value:  "postgres",
			EnvVar: "DB_DRIVER",
			Usage:  "database driver",
		},
		cli.StringFlag{
			Name:   "secret-key",
			EnvVar: "SECRET_KEY",
			Usage:  "secret key",
		},
		cli.StringFlag{
			Name:   "server-addr",
			Value:  ":8000",
			EnvVar: "SERVER_ADDR",
			Usage:  "server driver",
		},
	}
	app.Commands = []cli.Command{
		{
			Name:  "createadmin",
			Usage: "create an admin user",
			Action: func(c *cli.Context) error {

				// Create the database connection
				d, err := db.Connect(&db.Config{
					Driver: c.GlobalString("db-driver"),
					Args:   c.GlobalString("db-args"),
				})
				if err != nil {
					return err
				}
				defer d.Close()

				// Prompt for username
				var username string
				fmt.Print("Username? ")
				fmt.Scanln(&username)

				// Prompt for the password, hiding the input
				fmt.Print("Password? ")
				p, err := gopass.GetPasswd()
				if err != nil {
					return err
				}

				// Prompt for email address
				var email string
				fmt.Print("Email? ")
				fmt.Scanln(&email)

				// Create the new user
				u := &db.User{
					Username: username,
					Email:    email,
					IsAdmin:  true,
				}
				if err := u.SetPassword(string(p)); err != nil {
					return err
				}

				// Store the user in the database
				if err := d.C.Create(u).Error; err != nil {
					return err
				}

				return nil
			},
		},
		{
			Name:  "migrate",
			Usage: "perform all database migrations",
			Action: func(c *cli.Context) error {

				// Create the database connection
				d, err := db.Connect(&db.Config{
					Driver: c.GlobalString("db-driver"),
					Args:   c.GlobalString("db-args"),
				})
				if err != nil {
					return err
				}
				defer d.Close()

				// Perform all migrations
				return d.Migrate()
			},
		},
	}
	app.Action = func(c *cli.Context) error {

		// Create the database connection
		d, err := db.Connect(&db.Config{
			Driver: c.String("db-driver"),
			Args:   c.String("db-args"),
		})
		if err != nil {
			return err
		}
		defer d.Close()

		// Create the tweeter
		t := tweeter.New(&tweeter.Config{
			Database: d,
		})
		defer t.Close()

		// Create the server
		s, err := server.New(&server.Config{
			Addr:     c.String("server-addr"),
			Database: d,
			Tweeter:  t,
		})
		if err != nil {
			return err
		}
		defer s.Close()

		// Wait for SIGINT or SIGTERM
		sigChan := make(chan os.Signal)
		signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
		<-sigChan

		return nil
	}
	app.Run(os.Args)
}
<commit_msg>Add debug flag.<commit_after>package main

import (
	"fmt"
	"os"
	"os/signal"
	"syscall"

	"github.com/ReformedDevs/anonbot/db"
	"github.com/ReformedDevs/anonbot/server"
	"github.com/ReformedDevs/anonbot/tweeter"
	"github.com/howeyc/gopass"
	"github.com/sirupsen/logrus"
	"github.com/urfave/cli"
)

func main() {
	app := cli.NewApp()
	app.Name = "anonbot"
	app.Usage = "Twitter anon account bot"
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "db-args",
			Value:  "dbname=postgres user=postgres",
			EnvVar: "DB_ARGS",
			Usage:  "database arguments",
		},
		cli.StringFlag{
			Name:   "db-driver",
			Value:  "postgres",
			EnvVar: "DB_DRIVER",
			Usage:  "database driver",
		},
		cli.BoolFlag{
			Name:   "debug",
			EnvVar: "DEBUG",
			Usage:  "enable debug logging",
		},
		cli.StringFlag{
			Name:   "secret-key",
			EnvVar: "SECRET_KEY",
			Usage:  "secret key",
		},
		cli.StringFlag{
			Name:   "server-addr",
			Value:  ":8000",
			EnvVar: "SERVER_ADDR",
			Usage:  "server driver",
		},
	}
	app.Before = func(c *cli.Context) error {
		if c.Bool("debug") {
			logrus.SetLevel(logrus.DebugLevel)
		}
		return nil
	}
	app.Commands = []cli.Command{
		{
			Name:  "createadmin",
			Usage: "create an admin user",
			Action: func(c *cli.Context) error {

				// Create the database connection
				d, err := db.Connect(&db.Config{
					Driver: c.GlobalString("db-driver"),
					Args:   c.GlobalString("db-args"),
				})
				if err != nil {
					return err
				}
				defer d.Close()

				// Prompt for username
				var username string
				fmt.Print("Username? ")
				fmt.Scanln(&username)

				// Prompt for the password, hiding the input
				fmt.Print("Password? ")
				p, err := gopass.GetPasswd()
				if err != nil {
					return err
				}

				// Prompt for email address
				var email string
				fmt.Print("Email? ")
				fmt.Scanln(&email)

				// Create the new user
				u := &db.User{
					Username: username,
					Email:    email,
					IsAdmin:  true,
				}
				if err := u.SetPassword(string(p)); err != nil {
					return err
				}

				// Store the user in the database
				if err := d.C.Create(u).Error; err != nil {
					return err
				}

				return nil
			},
		},
		{
			Name:  "migrate",
			Usage: "perform all database migrations",
			Action: func(c *cli.Context) error {

				// Create the database connection
				d, err := db.Connect(&db.Config{
					Driver: c.GlobalString("db-driver"),
					Args:   c.GlobalString("db-args"),
				})
				if err != nil {
					return err
				}
				defer d.Close()

				// Perform all migrations
				return d.Migrate()
			},
		},
	}
	app.Action = func(c *cli.Context) error {

		// Create the database connection
		d, err := db.Connect(&db.Config{
			Driver: c.String("db-driver"),
			Args:   c.String("db-args"),
		})
		if err != nil {
			return err
		}
		defer d.Close()

		// Create the tweeter
		t := tweeter.New(&tweeter.Config{
			Database: d,
		})
		defer t.Close()

		// Create the server
		s, err := server.New(&server.Config{
			Addr:     c.String("server-addr"),
			Database: d,
			Tweeter:  t,
		})
		if err != nil {
			return err
		}
		defer s.Close()

		// Wait for SIGINT or SIGTERM
		sigChan := make(chan os.Signal)
		signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
		<-sigChan

		return nil
	}
	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"

	docopt "github.com/docopt/docopt-go"
	"github.com/luistm/go-bank-cli/cmd/bankcli/commands"
)

var intro = "    \n    Your command line finance manager.\n\n"

var usage = `Usage:
	bankcli -h | --help
	bankcli report --input <file> [ --grouped ]
	bankcli load --input <file>
	bankcli seller change <id> --pretty <name>
	bankcli seller new <name>
	bankcli seller show`

var options = `

Options:
	--grouped     The report is present grouped by seller
	--input       The path to the records list.
	--name        Specifies the name.
	-h --help     Show this screen.`

var version = "bankcli 1.0.0"

func main() {
	arguments, err := docopt.Parse(intro+usage+options, nil, true, version, false)
	if err != nil {
		fmt.Fprintln(os.Stderr, err.Error())
		os.Exit(2)
	}

	command, _ := commands.New(os.Args[1:])
	response := command.Execute(arguments)
	if response.String() != "" {
		fmt.Println(response.String())
		os.Exit(2)
	}
}
<commit_msg>Tags new version<commit_after>package main

import (
	"fmt"
	"os"

	docopt "github.com/docopt/docopt-go"
	"github.com/luistm/go-bank-cli/cmd/bankcli/commands"
)

var intro = "    \n    Your command line finance manager.\n\n"

var usage = `Usage:
	bankcli -h | --help
	bankcli report --input <file> [ --grouped ]
	bankcli load --input <file>
	bankcli seller change <id> --pretty <name>
	bankcli seller new <name>
	bankcli seller show`

var options = `

Options:
	--grouped     The report is present grouped by seller
	--input       The path to the records list.
	--name        Specifies the name.
	-h --help     Show this screen.`

var version = "bankcli 1.1.0"

func main() {
	arguments, err := docopt.Parse(intro+usage+options, nil, true, version, false)
	if err != nil {
		fmt.Fprintln(os.Stderr, err.Error())
		os.Exit(2)
	}

	command, _ := commands.New(os.Args[1:])
	response := command.Execute(arguments)
	if response.String() != "" {
		fmt.Println(response.String())
		os.Exit(2)
	}
}
<|endoftext|>


<commit_before>// Copyright 2021 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"log"
	"os"
	"strings"
)

func goTest(goroot string) error {
	log.Printf("Running Go test benchmarks for GOROOT %s", goroot)

	cmd := goCommand(goroot, "test", "-v", "-run=none", "-bench=.", "-count=5", "golang.org/x/benchmarks/...")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	env := os.Environ()
	needGOROOT := true
	for i := range env {
		if strings.HasPrefix(env[i], "GOROOT=") {
			env[i] = "GOROOT=" + goroot
			needGOROOT = false
		}
	}
	if needGOROOT {
		env = append(env, "GOROOT=" + goroot)
	}
	cmd.Env = env

	return cmd.Run()
}
<commit_msg>cmd/bench: go fmt<commit_after>// Copyright 2021 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"log"
	"os"
	"strings"
)

func goTest(goroot string) error {
	log.Printf("Running Go test benchmarks for GOROOT %s", goroot)

	cmd := goCommand(goroot, "test", "-v", "-run=none", "-bench=.", "-count=5", "golang.org/x/benchmarks/...")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	env := os.Environ()
	needGOROOT := true
	for i := range env {
		if strings.HasPrefix(env[i], "GOROOT=") {
			env[i] = "GOROOT=" + goroot
			needGOROOT = false
		}
	}
	if needGOROOT {
		env = append(env, "GOROOT="+goroot)
	}
	cmd.Env = env

	return cmd.Run()
}
<|endoftext|>


<commit_before>/*
 * MinIO Cloud Storage, (C) 2020 MinIO, 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 cmd

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"path"
	"sync"
	"time"

	"github.com/minio/minio/cmd/config"
	"github.com/minio/minio/cmd/logger"
	"github.com/minio/minio/pkg/env"
	"github.com/minio/minio/pkg/event"
	"github.com/minio/minio/pkg/madmin"
)

// BucketQuotaSys - map of bucket and quota configuration.
type BucketQuotaSys struct {
	sync.RWMutex
	quotaMap map[string]madmin.BucketQuota
}

// Set - set quota configuration.
func (sys *BucketQuotaSys) Set(bucketName string, q madmin.BucketQuota) {
	sys.Lock()
	sys.quotaMap[bucketName] = q
	sys.Unlock()
}

// Get - Get quota configuration.
func (sys *BucketQuotaSys) Get(bucketName string) (q madmin.BucketQuota, ok bool) {
	sys.RLock()
	defer sys.RUnlock()
	q, ok = sys.quotaMap[bucketName]
	return
}

// Remove - removes quota configuration.
func (sys *BucketQuotaSys) Remove(bucketName string) {
	sys.Lock()
	delete(sys.quotaMap, bucketName)
	sys.Unlock()
}

// Buckets - list of buckets with quota configuration
func (sys *BucketQuotaSys) Buckets() []string {
	sys.RLock()
	defer sys.RUnlock()
	var buckets []string
	for k := range sys.quotaMap {
		buckets = append(buckets, k)
	}
	return buckets
}

// Init initialize bucket quota sys configuration with all buckets.
func (sys *BucketQuotaSys) Init(buckets []BucketInfo, objAPI ObjectLayer) error {
	if objAPI == nil {
		return errServerNotInitialized
	}

	// In gateway mode, we do not support bucket quota.
	// So, this is a no-op for gateway servers.
	if globalIsGateway {
		return nil
	}

	return sys.load(buckets, objAPI)
}

func (sys *BucketQuotaSys) load(buckets []BucketInfo, objAPI ObjectLayer) error {
	for _, bucket := range buckets {
		ctx := logger.SetReqInfo(GlobalContext, &logger.ReqInfo{BucketName: bucket.Name})
		configFile := path.Join(bucketConfigPrefix, bucket.Name, bucketQuotaConfigFile)
		configData, err := readConfig(ctx, objAPI, configFile)
		if err != nil {
			if errors.Is(err, errConfigNotFound) {
				continue
			}
			return err
		}
		quotaCfg, err := parseBucketQuota(configData)
		if err != nil {
			return err
		}
		globalBucketQuotaSys.Set(bucket.Name, quotaCfg)
	}
	return nil
}

// NewBucketQuotaSys returns initialized BucketQuotaSys
func NewBucketQuotaSys() *BucketQuotaSys {
	return &BucketQuotaSys{quotaMap: map[string]madmin.BucketQuota{}}
}

// parseBucketQuota parses BucketQuota from json
func parseBucketQuota(data []byte) (quotaCfg madmin.BucketQuota, err error) {
	err = json.Unmarshal(data, &quotaCfg)
	if err != nil {
		return
	}
	if !quotaCfg.Type.IsValid() {
		return quotaCfg, fmt.Errorf("Invalid quota type %s", quotaCfg.Type)
	}
	return
}

type bucketStorageCache struct {
	bucketsSizes map[string]uint64
	lastUpdate   time.Time
	mu           sync.Mutex
}

func (b *bucketStorageCache) check(ctx context.Context, q madmin.BucketQuota, bucket string, size int64) error {
	b.mu.Lock()
	defer b.mu.Unlock()
	if time.Since(b.lastUpdate) > 10*time.Second {
		dui, err := loadDataUsageFromBackend(ctx, newObjectLayerWithoutSafeModeFn())
		if err != nil {
			return err
		}
		b.lastUpdate = time.Now()
		b.bucketsSizes = dui.BucketsSizes
	}
	currUsage := b.bucketsSizes[bucket]
	if (currUsage + uint64(size)) > q.Quota {
		return BucketQuotaExceeded{Bucket: bucket}
	}
	return nil
}
func enforceBucketQuota(ctx context.Context, bucket string, size int64) error {
	if size < 0 {
		return nil
	}
	q, ok := globalBucketQuotaSys.Get(bucket)
	if !ok {
		return nil
	}

	return globalBucketStorageCache.check(ctx, q, bucket, size)
}

const (
	bgQuotaInterval = 1 * time.Hour
)

// initQuotaEnforcement starts the routine that deletes objects in bucket
// that exceeds the FIFO quota
func initQuotaEnforcement(ctx context.Context, objAPI ObjectLayer) {
	if env.Get(envDataUsageCrawlConf, config.EnableOn) == config.EnableOn {
		go startBucketQuotaEnforcement(ctx, objAPI)
	}
}

func startBucketQuotaEnforcement(ctx context.Context, objAPI ObjectLayer) {
	for {
		select {
		case <-ctx.Done():
			return
		case <-time.NewTimer(bgQuotaInterval).C:
			logger.LogIf(ctx, enforceFIFOQuota(ctx, objAPI))
		}

	}
}

// enforceFIFOQuota deletes objects in FIFO order until sufficient objects
// have been deleted so as to bring bucket usage within quota
func enforceFIFOQuota(ctx context.Context, objectAPI ObjectLayer) error {
	// Turn off quota enforcement if data usage info is unavailable.
	if env.Get(envDataUsageCrawlConf, config.EnableOn) == config.EnableOff {
		return nil
	}
	for _, bucket := range globalBucketQuotaSys.Buckets() {
		// Check if the current bucket has quota restrictions, if not skip it
		cfg, ok := globalBucketQuotaSys.Get(bucket)
		if !ok {
			continue
		}
		if cfg.Type != madmin.FIFOQuota {
			continue
		}
		_, bucketHasLockConfig := globalBucketObjectLockSys.Get(bucket)

		dataUsageInfo, err := loadDataUsageFromBackend(ctx, objectAPI)
		if err != nil {
			return err
		}
		var toFree uint64
		if dataUsageInfo.BucketsSizes[bucket] > cfg.Quota {
			toFree = dataUsageInfo.BucketsSizes[bucket] - cfg.Quota
		}
		if toFree <= 0 {
			continue
		}
		// Allocate new results channel to receive ObjectInfo.
		objInfoCh := make(chan ObjectInfo)

		// Walk through all objects
		if err := objectAPI.Walk(ctx, bucket, "", objInfoCh); err != nil {
			return err
		}
		// reuse the fileScorer used by disk cache to score entries by
		// ModTime to find the oldest objects in bucket to delete. In
		// the context of bucket quota enforcement - number of hits are
		// irrelevant.
		scorer, err := newFileScorer(toFree, time.Now().Unix(), 1)
		if err != nil {
			return err
		}

		for obj := range objInfoCh {
			// skip objects currently under retention
			if bucketHasLockConfig && enforceRetentionForDeletion(ctx, obj) {
				continue
			}
			scorer.addFile(obj.Name, obj.ModTime, obj.Size, 1)
		}
		var objects []string
		numKeys := len(scorer.fileNames())
		for i, key := range scorer.fileNames() {
			objects = append(objects, key)
			if len(objects) < maxObjectList && (i < numKeys-1) {
				// skip deletion until maxObjectList or end of slice
				continue
			}

			if len(objects) == 0 {
				break
			}
			// Deletes a list of objects.
			deleteErrs, err := objectAPI.DeleteObjects(ctx, bucket, objects)
			if err != nil {
				logger.LogIf(ctx, err)
			} else {
				for i := range deleteErrs {
					if deleteErrs[i] != nil {
						logger.LogIf(ctx, deleteErrs[i])
						continue
					}
					// Notify object deleted event.
					sendEvent(eventArgs{
						EventName:  event.ObjectRemovedDelete,
						BucketName: bucket,
						Object: ObjectInfo{
							Name: objects[i],
						},
						Host: "Internal: [FIFO-QUOTA-EXPIRY]",
					})
				}
				objects = nil
			}
		}
	}
	return nil
}
<commit_msg>Allow writes for bucket exceeding FIFO quota (#9575)<commit_after>/*
 * MinIO Cloud Storage, (C) 2020 MinIO, 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 cmd

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"path"
	"sync"
	"time"

	"github.com/minio/minio/cmd/config"
	"github.com/minio/minio/cmd/logger"
	"github.com/minio/minio/pkg/env"
	"github.com/minio/minio/pkg/event"
	"github.com/minio/minio/pkg/madmin"
)

// BucketQuotaSys - map of bucket and quota configuration.
type BucketQuotaSys struct {
	sync.RWMutex
	quotaMap map[string]madmin.BucketQuota
}

// Set - set quota configuration.
func (sys *BucketQuotaSys) Set(bucketName string, q madmin.BucketQuota) {
	sys.Lock()
	sys.quotaMap[bucketName] = q
	sys.Unlock()
}

// Get - Get quota configuration.
func (sys *BucketQuotaSys) Get(bucketName string) (q madmin.BucketQuota, ok bool) {
	sys.RLock()
	defer sys.RUnlock()
	q, ok = sys.quotaMap[bucketName]
	return
}

// Remove - removes quota configuration.
func (sys *BucketQuotaSys) Remove(bucketName string) {
	sys.Lock()
	delete(sys.quotaMap, bucketName)
	sys.Unlock()
}

// Buckets - list of buckets with quota configuration
func (sys *BucketQuotaSys) Buckets() []string {
	sys.RLock()
	defer sys.RUnlock()
	var buckets []string
	for k := range sys.quotaMap {
		buckets = append(buckets, k)
	}
	return buckets
}

// Init initialize bucket quota sys configuration with all buckets.
func (sys *BucketQuotaSys) Init(buckets []BucketInfo, objAPI ObjectLayer) error {
	if objAPI == nil {
		return errServerNotInitialized
	}

	// In gateway mode, we do not support bucket quota.
	// So, this is a no-op for gateway servers.
	if globalIsGateway {
		return nil
	}

	return sys.load(buckets, objAPI)
}

func (sys *BucketQuotaSys) load(buckets []BucketInfo, objAPI ObjectLayer) error {
	for _, bucket := range buckets {
		ctx := logger.SetReqInfo(GlobalContext, &logger.ReqInfo{BucketName: bucket.Name})
		configFile := path.Join(bucketConfigPrefix, bucket.Name, bucketQuotaConfigFile)
		configData, err := readConfig(ctx, objAPI, configFile)
		if err != nil {
			if errors.Is(err, errConfigNotFound) {
				continue
			}
			return err
		}
		quotaCfg, err := parseBucketQuota(configData)
		if err != nil {
			return err
		}
		globalBucketQuotaSys.Set(bucket.Name, quotaCfg)
	}
	return nil
}

// NewBucketQuotaSys returns initialized BucketQuotaSys
func NewBucketQuotaSys() *BucketQuotaSys {
	return &BucketQuotaSys{quotaMap: map[string]madmin.BucketQuota{}}
}

// parseBucketQuota parses BucketQuota from json
func parseBucketQuota(data []byte) (quotaCfg madmin.BucketQuota, err error) {
	err = json.Unmarshal(data, &quotaCfg)
	if err != nil {
		return
	}
	if !quotaCfg.Type.IsValid() {
		return quotaCfg, fmt.Errorf("Invalid quota type %s", quotaCfg.Type)
	}
	return
}

type bucketStorageCache struct {
	bucketsSizes map[string]uint64
	lastUpdate   time.Time
	mu           sync.Mutex
}

func (b *bucketStorageCache) check(ctx context.Context, q madmin.BucketQuota, bucket string, size int64) error {
	b.mu.Lock()
	defer b.mu.Unlock()
	if time.Since(b.lastUpdate) > 10*time.Second {
		dui, err := loadDataUsageFromBackend(ctx, newObjectLayerWithoutSafeModeFn())
		if err != nil {
			return err
		}
		b.lastUpdate = time.Now()
		b.bucketsSizes = dui.BucketsSizes
	}
	currUsage := b.bucketsSizes[bucket]
	if (currUsage + uint64(size)) > q.Quota {
		return BucketQuotaExceeded{Bucket: bucket}
	}
	return nil
}
func enforceBucketQuota(ctx context.Context, bucket string, size int64) error {
	if size < 0 {
		return nil
	}
	q, ok := globalBucketQuotaSys.Get(bucket)
	if !ok {
		return nil
	}

	if q.Type == madmin.FIFOQuota {
		return nil
	}
	return globalBucketStorageCache.check(ctx, q, bucket, size)
}

const (
	bgQuotaInterval = 1 * time.Hour
)

// initQuotaEnforcement starts the routine that deletes objects in bucket
// that exceeds the FIFO quota
func initQuotaEnforcement(ctx context.Context, objAPI ObjectLayer) {
	if env.Get(envDataUsageCrawlConf, config.EnableOn) == config.EnableOn {
		go startBucketQuotaEnforcement(ctx, objAPI)
	}
}

func startBucketQuotaEnforcement(ctx context.Context, objAPI ObjectLayer) {
	for {
		select {
		case <-ctx.Done():
			return
		case <-time.NewTimer(bgQuotaInterval).C:
			logger.LogIf(ctx, enforceFIFOQuota(ctx, objAPI))
		}

	}
}

// enforceFIFOQuota deletes objects in FIFO order until sufficient objects
// have been deleted so as to bring bucket usage within quota
func enforceFIFOQuota(ctx context.Context, objectAPI ObjectLayer) error {
	// Turn off quota enforcement if data usage info is unavailable.
	if env.Get(envDataUsageCrawlConf, config.EnableOn) == config.EnableOff {
		return nil
	}
	for _, bucket := range globalBucketQuotaSys.Buckets() {
		// Check if the current bucket has quota restrictions, if not skip it
		cfg, ok := globalBucketQuotaSys.Get(bucket)
		if !ok {
			continue
		}
		if cfg.Type != madmin.FIFOQuota {
			continue
		}
		_, bucketHasLockConfig := globalBucketObjectLockSys.Get(bucket)

		dataUsageInfo, err := loadDataUsageFromBackend(ctx, objectAPI)
		if err != nil {
			return err
		}
		var toFree uint64
		if dataUsageInfo.BucketsSizes[bucket] > cfg.Quota {
			toFree = dataUsageInfo.BucketsSizes[bucket] - cfg.Quota
		}
		if toFree <= 0 {
			continue
		}
		// Allocate new results channel to receive ObjectInfo.
		objInfoCh := make(chan ObjectInfo)

		// Walk through all objects
		if err := objectAPI.Walk(ctx, bucket, "", objInfoCh); err != nil {
			return err
		}
		// reuse the fileScorer used by disk cache to score entries by
		// ModTime to find the oldest objects in bucket to delete. In
		// the context of bucket quota enforcement - number of hits are
		// irrelevant.
		scorer, err := newFileScorer(toFree, time.Now().Unix(), 1)
		if err != nil {
			return err
		}

		for obj := range objInfoCh {
			// skip objects currently under retention
			if bucketHasLockConfig && enforceRetentionForDeletion(ctx, obj) {
				continue
			}
			scorer.addFile(obj.Name, obj.ModTime, obj.Size, 1)
		}
		var objects []string
		numKeys := len(scorer.fileNames())
		for i, key := range scorer.fileNames() {
			objects = append(objects, key)
			if len(objects) < maxObjectList && (i < numKeys-1) {
				// skip deletion until maxObjectList or end of slice
				continue
			}

			if len(objects) == 0 {
				break
			}
			// Deletes a list of objects.
			deleteErrs, err := objectAPI.DeleteObjects(ctx, bucket, objects)
			if err != nil {
				logger.LogIf(ctx, err)
			} else {
				for i := range deleteErrs {
					if deleteErrs[i] != nil {
						logger.LogIf(ctx, deleteErrs[i])
						continue
					}
					// Notify object deleted event.
					sendEvent(eventArgs{
						EventName:  event.ObjectRemovedDelete,
						BucketName: bucket,
						Object: ObjectInfo{
							Name: objects[i],
						},
						Host: "Internal: [FIFO-QUOTA-EXPIRY]",
					})
				}
				objects = nil
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// Graceful restarts and shutdowns are not supported
// on windows, so using usual http.Serve instead.
//
// +build windows

package main

import (
	"net/http"
)

// serve is a wrapper on standard http.Serve method.
func serve(s *http.Server) error {
	return http.Serve(s)
}
<commit_msg>Fix Win's version of serve<commit_after>// Graceful restarts and shutdowns are not supported
// on windows, so using usual http.Serve instead.
//
// +build windows

package main

import (
	"net/http"
)

// serve is a wrapper on standard http.Serve method.
func serve(s *http.Server) error {
	return s.ListenAndServe()
}
<|endoftext|>