<commit_before>/* This file is part of VoltDB.
 * Copyright (C) 2008-2018 VoltDB Inc.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with VoltDB.  If not, see <http://www.gnu.org/licenses/>.
 */

package voltdbclient

import (
	"bytes"
	"database/sql/driver"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"time"

	"github.com/VoltDB/voltdb-client-go/wire"
)

// start back pressure when this many bytes are queued for write
const maxQueuedBytes = 262144
const maxResponseBuffer = 10000

type nodeConn struct {
	connInfo string
	connData *wire.ConnInfo
	tcpConn  *net.TCPConn

	drainCh chan chan bool
	bpCh    chan chan bool
	closeCh chan chan bool

	// channel for pi's meant specifically for this connection.
	ncPiCh  chan *procedureInvocation
	decoder *wire.Decoder
	encoder *wire.Encoder
}

func newNodeConn(ci string, ncPiCh chan *procedureInvocation) *nodeConn {
	return &nodeConn{
		connInfo: ci,
		ncPiCh:   ncPiCh,
		bpCh:     make(chan chan bool),
		closeCh:  make(chan chan bool),
		drainCh:  make(chan chan bool),
		decoder:  wire.NewDecoder(nil),
		encoder:  wire.NewEncoder(),
	}
}

func (nc *nodeConn) submit(pi *procedureInvocation) {
	nc.ncPiCh <- pi
}

// when the node conn is closed by its owning distributer
func (nc *nodeConn) close() chan bool {
	respCh := make(chan bool, 1)
	nc.closeCh <- respCh
	return respCh
}

func (nc *nodeConn) connect(protocolVersion int, piCh <-chan *procedureInvocation) error {
	tcpConn, connData, err := nc.networkConnect(protocolVersion)
	if err != nil {
		return err
	}
	nc.connData = connData
	nc.tcpConn = tcpConn

	responseCh := make(chan *bytes.Buffer, maxResponseBuffer)
	go nc.listen(responseCh)

	nc.drainCh = make(chan chan bool, 1)

	go nc.loop(piCh, responseCh, nc.bpCh, nc.drainCh)
	return nil
}

// called when the network listener loses connection.
// the 'processAsyncs' goroutine and channel stay in place over
// a reconnect, they're not affected.
func (nc *nodeConn) reconnect(protocolVersion int, piCh <-chan *procedureInvocation) {
	for {
		tcpConn, connData, err := nc.networkConnect(protocolVersion)
		if err != nil {
			log.Println(fmt.Printf("Failed to reconnect to server with %s, retrying\n", err))
			time.Sleep(5 * time.Second)
			continue
		}
		nc.tcpConn = tcpConn
		nc.connData = connData

		responseCh := make(chan *bytes.Buffer, maxResponseBuffer)
		go nc.listen(responseCh)
		go nc.loop(piCh, responseCh, nc.bpCh, nc.drainCh)
		break
	}
}

func (nc *nodeConn) networkConnect(protocolVersion int) (*net.TCPConn, *wire.ConnInfo, error) {
	defer func() {
		nc.decoder.Reset()
		nc.encoder.Reset()
	}()
	u, err := parseURL(nc.connInfo)
	if err != nil {
		return nil, nil, err
	}
	raddr, err := net.ResolveTCPAddr("tcp", u.Host)
	if err != nil {
		return nil, nil, fmt.Errorf("error resolving %v", nc.connInfo)
	}
	tcpConn, err := net.DialTCP("tcp", nil, raddr)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to connect to server %v", nc.connInfo)
	}
	pass, _ := u.User.Password()
	nc.encoder.Reset()
	login, err := nc.encoder.Login(protocolVersion, u.User.Username(), pass)
	if err != nil {
		tcpConn.Close()
		return nil, nil, fmt.Errorf("failed to serialize login message %v", nc.connInfo)
	}
	_, err = tcpConn.Write(login)
	if err != nil {
		return nil, nil, err
	}
	nc.decoder.Reset()
	nc.decoder.SetReader(tcpConn)
	i, err := nc.decoder.Login()
	if err != nil {
		tcpConn.Close()
		return nil, nil, fmt.Errorf("failed to login to server %v", nc.connInfo)
	}
	return tcpConn, i, nil
}

func (nc *nodeConn) drain(respCh chan bool) {
	nc.drainCh <- respCh
}

func (nc *nodeConn) hasBP() bool {
	respCh := make(chan bool)
	nc.bpCh <- respCh
	return <-respCh
}

// listen listens for messages from the server and calls back a registered listener.
// listen blocks on input from the server and should be run as a go routine.
func (nc *nodeConn) listen(responseCh chan<- *bytes.Buffer) {
	d := wire.NewDecoder(nc.tcpConn)
	s := &wire.Decoder{}
	for {
		b, err := d.Message()
		if err != nil {
			if responseCh == nil {
				// exiting
				return
			}
			// TODO: put the error on the channel
			// the owner needs to reconnect
			return
		}
		buf := bytes.NewBuffer(b)
		s.SetReader(buf)
		_, err = s.Byte()
		if err != nil {
			if responseCh == nil {
				return
			}
			return
		}
		responseCh <- buf
	}
}

func (nc *nodeConn) loop(piCh <-chan *procedureInvocation, responseCh <-chan *bytes.Buffer, bpCh <-chan chan bool, drainCh chan chan bool) {
	// declare mutable state
	requests := make(map[int64]*networkRequest)
	ncPiCh := nc.ncPiCh
	var draining bool
	var drainRespCh chan bool
	var queuedBytes int
	var bp bool

	var tci = int64(DefaultQueryTimeout / 10)                    // timeout check interval
	tcc := time.NewTimer(time.Duration(tci) * time.Nanosecond).C // timeout check timer channel

	// for ping
	var pingTimeout = 2 * time.Minute
	pingSentTime := time.Now()
	var pingOutstanding bool
	for {
		// setup select cases
		if draining {
			if queuedBytes == 0 && len(nc.ncPiCh) == 0 && len(piCh) == 0 {
				drainRespCh <- true
				drainRespCh = nil
				draining = false
			}
		}

		if queuedBytes > maxQueuedBytes && ncPiCh != nil {
			ncPiCh = nil
			bp = true
		} else if ncPiCh == nil {
			ncPiCh = nc.ncPiCh
			bp = false
		}

		// ping
		pingSinceSent := time.Now().Sub(pingSentTime)
		if pingOutstanding {
			if pingSinceSent > pingTimeout {
				// TODO: should disconnect
			}
		} else if pingSinceSent > pingTimeout/3 {
			nc.sendPing(nc.tcpConn)
			pingOutstanding = true
			pingSentTime = time.Now()
		}

		select {
		case respCh := <-nc.closeCh:
			nc.tcpConn.Close()
			respCh <- true
			return
		case pi := <-ncPiCh:
			nc.handleProcedureInvocation(pi, &requests, &queuedBytes)
		case pi := <-piCh:
			nc.handleProcedureInvocation(pi, &requests, &queuedBytes)
		case resp := <-responseCh:
			nc.decoder.SetReader(resp)
			handle, err := nc.decoder.Int64()
			nc.decoder.Reset()
			// can't do anything without a handle.  If reading the handle fails,
			// then log and drop the message.
			if err != nil {
				continue
			}
			if handle == PingHandle {
				pingOutstanding = false
				continue
			}
			req := requests[handle]
			if req == nil {
				// there's a race here with timeout.  A request can be timed out and
				// then a response received.  In this case drop the response.
				continue
			}
			queuedBytes -= req.numBytes

			delete(requests, handle)
			if req.isSync() {
				nc.handleSyncResponse(handle, resp, req)
			} else {
				nc.handleAsyncResponse(handle, resp, req)
			}

		case respBPCh := <-bpCh:
			respBPCh <- bp
		case drainRespCh = <-drainCh:
			draining = true
		// check for timed out procedure invocations
		case <-tcc:
			for _, req := range requests {
				if time.Now().After(req.submitted.Add(req.timeout)) {
					queuedBytes -= req.numBytes
					nc.handleTimeout(req)
					delete(requests, req.handle)
				}
			}
			tcc = time.NewTimer(time.Duration(tci) * time.Nanosecond).C
		}
	}
}

func (nc *nodeConn) handleProcedureInvocation(pi *procedureInvocation, requests *map[int64]*networkRequest, queuedBytes *int) {
	var nr *networkRequest
	if pi.isAsync() {
		nr = newAsyncRequest(pi.handle, pi.responseCh, pi.isQuery, pi.arc, pi.getLen(), pi.timeout, time.Now())
	} else {
		nr = newSyncRequest(pi.handle, pi.responseCh, pi.isQuery, pi.getLen(), pi.timeout, time.Now())
	}
	(*requests)[pi.handle] = nr
	*queuedBytes += pi.slen
	nc.encoder.Reset()
	EncodePI(nc.encoder, pi)
	nc.tcpConn.Write(nc.encoder.Bytes())
	nc.encoder.Reset()
}

func (nc *nodeConn) handleSyncResponse(handle int64, r io.Reader, req *networkRequest) {
	respCh := req.getChan()
	nc.decoder.SetReader(r)
	defer nc.decoder.Reset()
	rsp, err := decodeResponse(nc.decoder, handle)
	if err != nil {
		respCh <- err.(voltResponse)
	} else if req.isQuery() {

		if rows, err := decodeRows(nc.decoder, rsp); err != nil {
			respCh <- err.(voltResponse)
		} else {
			respCh <- rows
		}
	} else {

		if result, err := decodeResult(nc.decoder, rsp); err != nil {
			respCh <- err.(voltResponse)
		} else {
			respCh <- result
		}
	}

}

func (nc *nodeConn) handleAsyncResponse(handle int64, r io.Reader, req *networkRequest) {
	d := wire.NewDecoder(r)
	rsp, err := decodeResponse(d, handle)
	if err != nil {
		req.arc.ConsumeError(err)
	} else if req.isQuery() {
		if rows, err := decodeRows(d, rsp); err != nil {
			req.arc.ConsumeError(err)
		} else {
			req.arc.ConsumeRows(rows)
		}
	} else {
		if result, err := decodeResult(d, rsp); err != nil {
			req.arc.ConsumeError(err)
		} else {
			req.arc.ConsumeResult(result)
		}
	}
}

func (nc *nodeConn) handleTimeout(req *networkRequest) {
	err := errors.New("timeout")
	verr := VoltError{voltResponse: emptyVoltResponseInfo(), error: err}
	if req.isSync() {
		respCh := req.getChan()
		respCh <- verr
	} else {
		req.arc.ConsumeError(verr)
	}
}

func (nc *nodeConn) sendPing(writer io.Writer) {
	pi := newProcedureInvocationByHandle(PingHandle, true, "@Ping", []driver.Value{})
	nc.encoder.Reset()
	EncodePI(nc.encoder, pi)
	writer.Write(nc.encoder.Bytes())
	nc.encoder.Reset()
}

// AsyncResponseConsumer is a type that consumes responses from asynchronous
// Queries and Execs.
// In the VoltDB go client, asynchronous requests are continuously processed by
// one or more goroutines executing in the background.  When a response from
// the server is received for an asynchronous request, one of the methods in
// this interface is invoked.  An instance of AyncResponseConsumer is passed
// when an asynchronous request is made, this instance will process the
// response for that request.
type AsyncResponseConsumer interface {

	// This method is invoked when an error is returned by an async Query
	// or an Exec.
	ConsumeError(error)
	// This method is invoked when a Result is returned by an async Exec.
	ConsumeResult(driver.Result)
	// This method is invoked when Rows is returned by an async Query.
	ConsumeRows(driver.Rows)
}

// Null Value type
type nullValue struct {
	colType int8
}

func (nv *nullValue) getColType() int8 {
	return nv.colType
}
<commit_msg>don't pass responseCh around<commit_after>/* This file is part of VoltDB.
 * Copyright (C) 2008-2018 VoltDB Inc.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with VoltDB.  If not, see <http://www.gnu.org/licenses/>.
 */

package voltdbclient

import (
	"bytes"
	"database/sql/driver"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"time"

	"github.com/VoltDB/voltdb-client-go/wire"
)

// start back pressure when this many bytes are queued for write
const maxQueuedBytes = 262144
const maxResponseBuffer = 10000

type nodeConn struct {
	connInfo string
	connData *wire.ConnInfo
	tcpConn  *net.TCPConn

	drainCh chan chan bool
	bpCh    chan chan bool
	closeCh chan chan bool

	// channel for pi's meant specifically for this connection.
	ncPiCh     chan *procedureInvocation
	decoder    *wire.Decoder
	encoder    *wire.Encoder
	responseCh chan *bytes.Buffer
}

func newNodeConn(ci string, ncPiCh chan *procedureInvocation) *nodeConn {
	return &nodeConn{
		connInfo:   ci,
		ncPiCh:     ncPiCh,
		bpCh:       make(chan chan bool),
		closeCh:    make(chan chan bool),
		drainCh:    make(chan chan bool),
		decoder:    wire.NewDecoder(nil),
		encoder:    wire.NewEncoder(),
		responseCh: make(chan *bytes.Buffer, maxResponseBuffer),
	}
}

func (nc *nodeConn) submit(pi *procedureInvocation) {
	nc.ncPiCh <- pi
}

// when the node conn is closed by its owning distributer
func (nc *nodeConn) close() chan bool {
	respCh := make(chan bool, 1)
	nc.closeCh <- respCh
	return respCh
}

func (nc *nodeConn) connect(protocolVersion int, piCh <-chan *procedureInvocation) error {
	tcpConn, connData, err := nc.networkConnect(protocolVersion)
	if err != nil {
		return err
	}
	nc.connData = connData
	nc.tcpConn = tcpConn

	go nc.listen()

	nc.drainCh = make(chan chan bool, 1)

	go nc.loop(piCh, nc.bpCh, nc.drainCh)
	return nil
}

// called when the network listener loses connection.
// the 'processAsyncs' goroutine and channel stay in place over
// a reconnect, they're not affected.
func (nc *nodeConn) reconnect(protocolVersion int, piCh <-chan *procedureInvocation) {
	for {
		tcpConn, connData, err := nc.networkConnect(protocolVersion)
		if err != nil {
			log.Println(fmt.Printf("Failed to reconnect to server with %s, retrying\n", err))
			time.Sleep(5 * time.Second)
			continue
		}
		nc.tcpConn = tcpConn
		nc.connData = connData
		go nc.listen()
		go nc.loop(piCh, nc.bpCh, nc.drainCh)
		break
	}
}

func (nc *nodeConn) networkConnect(protocolVersion int) (*net.TCPConn, *wire.ConnInfo, error) {
	defer func() {
		nc.decoder.Reset()
		nc.encoder.Reset()
	}()
	u, err := parseURL(nc.connInfo)
	if err != nil {
		return nil, nil, err
	}
	raddr, err := net.ResolveTCPAddr("tcp", u.Host)
	if err != nil {
		return nil, nil, fmt.Errorf("error resolving %v", nc.connInfo)
	}
	tcpConn, err := net.DialTCP("tcp", nil, raddr)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to connect to server %v", nc.connInfo)
	}
	pass, _ := u.User.Password()
	nc.encoder.Reset()
	login, err := nc.encoder.Login(protocolVersion, u.User.Username(), pass)
	if err != nil {
		tcpConn.Close()
		return nil, nil, fmt.Errorf("failed to serialize login message %v", nc.connInfo)
	}
	_, err = tcpConn.Write(login)
	if err != nil {
		return nil, nil, err
	}
	nc.decoder.Reset()
	nc.decoder.SetReader(tcpConn)
	i, err := nc.decoder.Login()
	if err != nil {
		tcpConn.Close()
		return nil, nil, fmt.Errorf("failed to login to server %v", nc.connInfo)
	}
	return tcpConn, i, nil
}

func (nc *nodeConn) drain(respCh chan bool) {
	nc.drainCh <- respCh
}

func (nc *nodeConn) hasBP() bool {
	respCh := make(chan bool)
	nc.bpCh <- respCh
	return <-respCh
}

// listen listens for messages from the server and calls back a registered listener.
// listen blocks on input from the server and should be run as a go routine.
func (nc *nodeConn) listen() {
	d := wire.NewDecoder(nc.tcpConn)
	s := &wire.Decoder{}
	for {
		b, err := d.Message()
		if err != nil {
			if nc.responseCh == nil {
				// exiting
				return
			}
			// TODO: put the error on the channel
			// the owner needs to reconnect
			return
		}
		buf := bytes.NewBuffer(b)
		s.SetReader(buf)
		_, err = s.Byte()
		if err != nil {
			if nc.responseCh == nil {
				return
			}
			return
		}
		nc.responseCh <- buf
	}
}

func (nc *nodeConn) loop(piCh <-chan *procedureInvocation, bpCh <-chan chan bool, drainCh chan chan bool) {
	// declare mutable state
	requests := make(map[int64]*networkRequest)
	ncPiCh := nc.ncPiCh
	var draining bool
	var drainRespCh chan bool
	var queuedBytes int
	var bp bool

	var tci = int64(DefaultQueryTimeout / 10)                    // timeout check interval
	tcc := time.NewTimer(time.Duration(tci) * time.Nanosecond).C // timeout check timer channel

	// for ping
	var pingTimeout = 2 * time.Minute
	pingSentTime := time.Now()
	var pingOutstanding bool
	for {
		// setup select cases
		if draining {
			if queuedBytes == 0 && len(nc.ncPiCh) == 0 && len(piCh) == 0 {
				drainRespCh <- true
				drainRespCh = nil
				draining = false
			}
		}

		if queuedBytes > maxQueuedBytes && ncPiCh != nil {
			ncPiCh = nil
			bp = true
		} else if ncPiCh == nil {
			ncPiCh = nc.ncPiCh
			bp = false
		}

		// ping
		pingSinceSent := time.Now().Sub(pingSentTime)
		if pingOutstanding {
			if pingSinceSent > pingTimeout {
				// TODO: should disconnect
			}
		} else if pingSinceSent > pingTimeout/3 {
			nc.sendPing(nc.tcpConn)
			pingOutstanding = true
			pingSentTime = time.Now()
		}

		select {
		case respCh := <-nc.closeCh:
			nc.tcpConn.Close()
			respCh <- true
			return
		case pi := <-ncPiCh:
			nc.handleProcedureInvocation(pi, &requests, &queuedBytes)
		case pi := <-piCh:
			nc.handleProcedureInvocation(pi, &requests, &queuedBytes)
		case resp := <-nc.responseCh:
			nc.decoder.SetReader(resp)
			handle, err := nc.decoder.Int64()
			nc.decoder.Reset()
			// can't do anything without a handle.  If reading the handle fails,
			// then log and drop the message.
			if err != nil {
				continue
			}
			if handle == PingHandle {
				pingOutstanding = false
				continue
			}
			req := requests[handle]
			if req == nil {
				// there's a race here with timeout.  A request can be timed out and
				// then a response received.  In this case drop the response.
				continue
			}
			queuedBytes -= req.numBytes

			delete(requests, handle)
			if req.isSync() {
				nc.handleSyncResponse(handle, resp, req)
			} else {
				nc.handleAsyncResponse(handle, resp, req)
			}

		case respBPCh := <-bpCh:
			respBPCh <- bp
		case drainRespCh = <-drainCh:
			draining = true
		// check for timed out procedure invocations
		case <-tcc:
			for _, req := range requests {
				if time.Now().After(req.submitted.Add(req.timeout)) {
					queuedBytes -= req.numBytes
					nc.handleTimeout(req)
					delete(requests, req.handle)
				}
			}
			tcc = time.NewTimer(time.Duration(tci) * time.Nanosecond).C
		}
	}
}

func (nc *nodeConn) handleProcedureInvocation(pi *procedureInvocation, requests *map[int64]*networkRequest, queuedBytes *int) {
	var nr *networkRequest
	if pi.isAsync() {
		nr = newAsyncRequest(pi.handle, pi.responseCh, pi.isQuery, pi.arc, pi.getLen(), pi.timeout, time.Now())
	} else {
		nr = newSyncRequest(pi.handle, pi.responseCh, pi.isQuery, pi.getLen(), pi.timeout, time.Now())
	}
	(*requests)[pi.handle] = nr
	*queuedBytes += pi.slen
	nc.encoder.Reset()
	EncodePI(nc.encoder, pi)
	nc.tcpConn.Write(nc.encoder.Bytes())
	nc.encoder.Reset()
}

func (nc *nodeConn) handleSyncResponse(handle int64, r io.Reader, req *networkRequest) {
	respCh := req.getChan()
	nc.decoder.SetReader(r)
	defer nc.decoder.Reset()
	rsp, err := decodeResponse(nc.decoder, handle)
	if err != nil {
		respCh <- err.(voltResponse)
	} else if req.isQuery() {

		if rows, err := decodeRows(nc.decoder, rsp); err != nil {
			respCh <- err.(voltResponse)
		} else {
			respCh <- rows
		}
	} else {

		if result, err := decodeResult(nc.decoder, rsp); err != nil {
			respCh <- err.(voltResponse)
		} else {
			respCh <- result
		}
	}

}

func (nc *nodeConn) handleAsyncResponse(handle int64, r io.Reader, req *networkRequest) {
	d := wire.NewDecoder(r)
	rsp, err := decodeResponse(d, handle)
	if err != nil {
		req.arc.ConsumeError(err)
	} else if req.isQuery() {
		if rows, err := decodeRows(d, rsp); err != nil {
			req.arc.ConsumeError(err)
		} else {
			req.arc.ConsumeRows(rows)
		}
	} else {
		if result, err := decodeResult(d, rsp); err != nil {
			req.arc.ConsumeError(err)
		} else {
			req.arc.ConsumeResult(result)
		}
	}
}

func (nc *nodeConn) handleTimeout(req *networkRequest) {
	err := errors.New("timeout")
	verr := VoltError{voltResponse: emptyVoltResponseInfo(), error: err}
	if req.isSync() {
		respCh := req.getChan()
		respCh <- verr
	} else {
		req.arc.ConsumeError(verr)
	}
}

func (nc *nodeConn) sendPing(writer io.Writer) {
	pi := newProcedureInvocationByHandle(PingHandle, true, "@Ping", []driver.Value{})
	nc.encoder.Reset()
	EncodePI(nc.encoder, pi)
	writer.Write(nc.encoder.Bytes())
	nc.encoder.Reset()
}

// AsyncResponseConsumer is a type that consumes responses from asynchronous
// Queries and Execs.
// In the VoltDB go client, asynchronous requests are continuously processed by
// one or more goroutines executing in the background.  When a response from
// the server is received for an asynchronous request, one of the methods in
// this interface is invoked.  An instance of AyncResponseConsumer is passed
// when an asynchronous request is made, this instance will process the
// response for that request.
type AsyncResponseConsumer interface {

	// This method is invoked when an error is returned by an async Query
	// or an Exec.
	ConsumeError(error)
	// This method is invoked when a Result is returned by an async Exec.
	ConsumeResult(driver.Result)
	// This method is invoked when Rows is returned by an async Query.
	ConsumeRows(driver.Rows)
}

// Null Value type
type nullValue struct {
	colType int8
}

func (nv *nullValue) getColType() int8 {
	return nv.colType
}
<|endoftext|>


<commit_before><commit_msg>update uuid path<commit_after><|endoftext|>


<commit_before>package tlc_test

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"testing"

	"github.com/itchio/wharf/tlc"
	"github.com/stretchr/testify/assert"
)

func Test_Walk(t *testing.T) {
	tmpPath := mktestdir(t, "walk")
	defer func() {
		err := os.RemoveAll(tmpPath)
		must(t, err)
	}()

	info, err := tlc.Walk(tmpPath, nil)
	must(t, err)

	dirs := []string{
		".",
		"foo",
		"foo/dir_a",
		"foo/dir_b",
	}
	for i, dir := range dirs {
		assert.Equal(t, dir, info.Dirs[i].Path, "dirs should be all listed")
	}

	files := []string{
		"foo/dir_a/baz",
		"foo/dir_a/bazzz",
		"foo/dir_b/zoom",
		"foo/file_f",
		"foo/file_z",
	}
	for i, file := range files {
		assert.Equal(t, file, info.Files[i].Path, "files should be all listed")
	}

	if testSymlinks {
		for i, symlink := range symlinks {
			assert.Equal(t, symlink.Newname, info.Symlinks[i].Path, "symlink should be at correct path")
			assert.Equal(t, symlink.Oldname, info.Symlinks[i].Dest, "symlink should point to correct path")
		}
	}
}

func Test_Prepare(t *testing.T) {
	tmpPath := mktestdir(t, "prepare")
	defer func() {
		err := os.RemoveAll(tmpPath)
		must(t, err)
	}()

	info, err := tlc.Walk(tmpPath, nil)
	must(t, err)

	tmpPath2, err := ioutil.TempDir(".", "tmp_prepare")
	must(t, err)

	err = info.Prepare(tmpPath2)
	must(t, err)

	info2, err := tlc.Walk(tmpPath2, nil)
	must(t, err)
	assert.Equal(t, info, info2, "must recreate same structure")
}

// Support code

func must(t *testing.T, err error) {
	if err != nil {
		t.Error("must failed: ", err.Error())
		t.FailNow()
	}
}

type regEntry struct {
	Path string
	Size int
	Byte byte
}

type symlinkEntry struct {
	Oldname string
	Newname string
}

var regulars = []regEntry{
	{"foo/file_f", 50, 0xd},
	{"foo/dir_a/baz", 10, 0xa},
	{"foo/dir_b/zoom", 30, 0xc},
	{"foo/file_z", 40, 0xe},
	{"foo/dir_a/bazzz", 20, 0xb},
}

var symlinks = []symlinkEntry{
	{"file_z", "foo/file_m"},
	{"dir_a/baz", "foo/file_o"},
}

var testSymlinks = runtime.GOOS != "windows"

func mktestdir(t *testing.T, name string) string {
	tmpPath, err := ioutil.TempDir(".", "tmp_"+name)
	must(t, err)

	must(t, os.RemoveAll(tmpPath))

	for _, entry := range regulars {
		fullPath := filepath.Join(tmpPath, entry.Path)
		must(t, os.MkdirAll(filepath.Dir(fullPath), os.FileMode(0777)))
		file, err := os.Create(fullPath)
		must(t, err)

		filler := []byte{entry.Byte}
		for i := 0; i < entry.Size; i++ {
			_, err := file.Write(filler)
			must(t, err)
		}
		must(t, file.Close())
	}

	if testSymlinks {
		for _, entry := range symlinks {
			new := filepath.Join(tmpPath, entry.Newname)
			must(t, os.Symlink(entry.Oldname, new))
		}
	}

	return tmpPath
}
<commit_msg>Fix tests<commit_after>package tlc_test

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"testing"

	"github.com/itchio/wharf/tlc"
	"github.com/stretchr/testify/assert"
)

func Test_Walk(t *testing.T) {
	tmpPath := mktestdir(t, "walk")
	defer func() {
		err := os.RemoveAll(tmpPath)
		must(t, err)
	}()

	info, err := tlc.Walk(tmpPath, nil)
	must(t, err)

	dirs := []string{
		"foo",
		"foo/dir_a",
		"foo/dir_b",
	}
	for i, dir := range dirs {
		assert.Equal(t, dir, info.Dirs[i].Path, "dirs should be all listed")
	}

	files := []string{
		"foo/dir_a/baz",
		"foo/dir_a/bazzz",
		"foo/dir_b/zoom",
		"foo/file_f",
		"foo/file_z",
	}
	for i, file := range files {
		assert.Equal(t, file, info.Files[i].Path, "files should be all listed")
	}

	if testSymlinks {
		for i, symlink := range symlinks {
			assert.Equal(t, symlink.Newname, info.Symlinks[i].Path, "symlink should be at correct path")
			assert.Equal(t, symlink.Oldname, info.Symlinks[i].Dest, "symlink should point to correct path")
		}
	}
}

func Test_Prepare(t *testing.T) {
	tmpPath := mktestdir(t, "prepare")
	defer func() {
		err := os.RemoveAll(tmpPath)
		must(t, err)
	}()

	info, err := tlc.Walk(tmpPath, nil)
	must(t, err)

	tmpPath2, err := ioutil.TempDir(".", "tmp_prepare")
	must(t, err)

	err = info.Prepare(tmpPath2)
	must(t, err)

	info2, err := tlc.Walk(tmpPath2, nil)
	must(t, err)
	assert.Equal(t, info, info2, "must recreate same structure")
}

// Support code

func must(t *testing.T, err error) {
	if err != nil {
		t.Error("must failed: ", err.Error())
		t.FailNow()
	}
}

type regEntry struct {
	Path string
	Size int
	Byte byte
}

type symlinkEntry struct {
	Oldname string
	Newname string
}

var regulars = []regEntry{
	{"foo/file_f", 50, 0xd},
	{"foo/dir_a/baz", 10, 0xa},
	{"foo/dir_b/zoom", 30, 0xc},
	{"foo/file_z", 40, 0xe},
	{"foo/dir_a/bazzz", 20, 0xb},
}

var symlinks = []symlinkEntry{
	{"file_z", "foo/file_m"},
	{"dir_a/baz", "foo/file_o"},
}

var testSymlinks = runtime.GOOS != "windows"

func mktestdir(t *testing.T, name string) string {
	tmpPath, err := ioutil.TempDir(".", "tmp_"+name)
	must(t, err)

	must(t, os.RemoveAll(tmpPath))

	for _, entry := range regulars {
		fullPath := filepath.Join(tmpPath, entry.Path)
		must(t, os.MkdirAll(filepath.Dir(fullPath), os.FileMode(0777)))
		file, err := os.Create(fullPath)
		must(t, err)

		filler := []byte{entry.Byte}
		for i := 0; i < entry.Size; i++ {
			_, err := file.Write(filler)
			must(t, err)
		}
		must(t, file.Close())
	}

	if testSymlinks {
		for _, entry := range symlinks {
			new := filepath.Join(tmpPath, entry.Newname)
			must(t, os.Symlink(entry.Oldname, new))
		}
	}

	return tmpPath
}
<|endoftext|>


<commit_before>package app

import (
	"fmt"

	"github.com/moncho/dry/appui"
	"github.com/moncho/dry/docker"
	"github.com/moncho/dry/ui"
	"github.com/nsf/termbox-go"
)

type nodesScreenEventHandler struct {
	baseEventHandler
	passingEvents bool
}

func (h *nodesScreenEventHandler) widget() appui.EventableWidget {
	return h.dry.state.activeWidget
}

func (h *nodesScreenEventHandler) handle(event termbox.Event) {
	if h.passingEvents {
		h.eventChan <- event
		return
	}
	handled := false
	focus := true

	switch event.Key {
	case termbox.KeyCtrlA:
		dry := h.dry
		rw := appui.NewAskForConfirmation("Changing node availability, please type one of ('active'|'pause'|'drain')")
		h.passingEvents = true
		handled = true
		dry.widgetRegistry.add(rw)
		go func() {
			events := ui.EventSource{
				Events: h.eventChan,
				EventHandledCallback: func(e termbox.Event) error {
					return refreshScreen()
				},
			}
			rw.OnFocus(events)
			dry.widgetRegistry.remove(rw)
			availability, canceled := rw.Text()
			h.passingEvents = false
			if canceled {
				return
			}
			if availability != "active" && availability != "pause" && availability != "drain" {
				dry.appmessage(fmt.Sprintf("Invalid availability: %s", availability))
				return
			}

			changeNode := func(nodeID string) error {
				err := dry.dockerDaemon.NodeChangeAvailabiliy(
					nodeID,
					docker.NewNodeAvailability(availability))

				if err == nil {
					dry.appmessage(fmt.Sprintf("Node %s availability is now %s", nodeID, availability))
				} else {
					dry.appmessage(fmt.Sprintf("Could not change node availability, error %s", err.Error()))
					return err
				}
				return refreshScreen()
			}
			h.widget().OnEvent(changeNode)
		}()

	case termbox.KeyEnter:
		showServices := func(nodeID string) error {
			h.dry.ShowTasks(nodeID)
			return refreshScreen()
		}
		h.widget().OnEvent(showServices)
		handled = true

	}
	if !handled {
		h.baseEventHandler.handle(event)
	} else {
		h.setFocus(focus)
		if h.hasFocus() {
			refreshScreen()
		}
	}
}

type taskScreenEventHandler struct {
	baseEventHandler
}

func (h *taskScreenEventHandler) handle(event termbox.Event) {

	switch event.Key {
	case termbox.KeyEsc:
		h.dry.ShowNodes()
	}

	h.baseEventHandler.handle(event)

}
<commit_msg>Link node widget with its events<commit_after>package app

import (
	"fmt"

	"github.com/moncho/dry/appui"
	"github.com/moncho/dry/docker"
	"github.com/moncho/dry/ui"
	"github.com/nsf/termbox-go"
)

type nodesScreenEventHandler struct {
	baseEventHandler
	passingEvents bool
}

func (h *nodesScreenEventHandler) widget() appui.EventableWidget {
	return h.dry.widgetRegistry.Nodes
}

func (h *nodesScreenEventHandler) handle(event termbox.Event) {
	if h.passingEvents {
		h.eventChan <- event
		return
	}
	handled := false
	focus := true

	switch event.Key {
	case termbox.KeyCtrlA:
		dry := h.dry
		rw := appui.NewAskForConfirmation("Changing node availability, please type one of ('active'|'pause'|'drain')")
		h.passingEvents = true
		handled = true
		dry.widgetRegistry.add(rw)
		go func() {
			events := ui.EventSource{
				Events: h.eventChan,
				EventHandledCallback: func(e termbox.Event) error {
					return refreshScreen()
				},
			}
			rw.OnFocus(events)
			dry.widgetRegistry.remove(rw)
			availability, canceled := rw.Text()
			h.passingEvents = false
			if canceled {
				return
			}
			if availability != "active" && availability != "pause" && availability != "drain" {
				dry.appmessage(fmt.Sprintf("Invalid availability: %s", availability))
				return
			}

			changeNode := func(nodeID string) error {
				err := dry.dockerDaemon.NodeChangeAvailabiliy(
					nodeID,
					docker.NewNodeAvailability(availability))

				if err == nil {
					dry.appmessage(fmt.Sprintf("Node %s availability is now %s", nodeID, availability))
				} else {
					dry.appmessage(fmt.Sprintf("Could not change node availability, error %s", err.Error()))
					return err
				}
				return refreshScreen()
			}
			h.widget().OnEvent(changeNode)
		}()

	case termbox.KeyEnter:
		showServices := func(nodeID string) error {
			h.dry.ShowTasks(nodeID)
			return refreshScreen()
		}
		h.widget().OnEvent(showServices)
		handled = true

	}
	if !handled {
		h.baseEventHandler.handle(event)
	} else {
		h.setFocus(focus)
		if h.hasFocus() {
			refreshScreen()
		}
	}
}

type taskScreenEventHandler struct {
	baseEventHandler
}

func (h *taskScreenEventHandler) handle(event termbox.Event) {

	switch event.Key {
	case termbox.KeyEsc:
		h.dry.ShowNodes()
	}

	h.baseEventHandler.handle(event)

}
<|endoftext|>


<commit_before>// Copyright (c) 2016-2019 Uber Technologies, 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 dockerregistry

import (
	"fmt"
	"io"

	"github.com/uber/kraken/lib/dockerregistry/transfer"
	"github.com/uber/kraken/lib/store"
	"github.com/uber/kraken/utils/log"

	"github.com/docker/distribution/context"
	storagedriver "github.com/docker/distribution/registry/storage/driver"
	"github.com/docker/distribution/registry/storage/driver/factory"
	"github.com/uber-go/tally"
)

// The path layout in the storage backend is roughly as follows:
//
//		<root>/v2
//			-> repositories/
// 				-><name>/
// 					-> _manifests/
// 						revisions
//							-> <manifest digest path>
//								-> link
// 						tags/<tag>
//							-> current/link
// 							-> index
//								-> <algorithm>/<hex digest>/link
// 					-> _layers/
// 						<layer links to blob store>
// 					-> _uploads/<id>
// 						data
// 						startedat
// 						hashstates/<algorithm>/<offset>
//			-> blobs/<algorithm>
//				<split directory content addressable storage>
//

const (
	// Name of storage driver
	Name            = "kraken"
	retries         = 3
	downloadTimeout = 120     //seconds
	readtimeout     = 15 * 60 //seconds
	writetimeout    = 15 * 60 //seconds
)

func init() {
	factory.Register(Name, &krakenStorageDriverFactory{})
}

// InvalidRequestError implements error and contains the path that is not supported
type InvalidRequestError struct {
	path string
}

func (e InvalidRequestError) Error() string {
	return fmt.Sprintf("invalid request: %s", e.path)
}

type krakenStorageDriverFactory struct{}

func getParam(params map[string]interface{}, name string) interface{} {
	p, ok := params[name]
	if !ok || p == nil {
		log.Fatalf("Required parameter %s not found", name)
	}
	return p
}

func (factory *krakenStorageDriverFactory) Create(
	params map[string]interface{}) (storagedriver.StorageDriver, error) {

	// Common parameters.
	constructor := getParam(params, "constructor").(string)
	config := getParam(params, "config").(Config)
	transferer := getParam(params, "transferer").(transfer.ImageTransferer)
	metrics := getParam(params, "metrics").(tally.Scope)

	switch constructor {
	case _rw:
		castore := getParam(params, "castore").(*store.CAStore)
		return NewReadWriteStorageDriver(config, castore, transferer, metrics), nil
	case _ro:
		blobstore := getParam(params, "blobstore").(BlobStore)
		return NewReadOnlyStorageDriver(config, blobstore, transferer, metrics), nil
	default:
		return nil, fmt.Errorf("unknown constructor %s", constructor)
	}
}

// KrakenStorageDriver is a storage driver
type KrakenStorageDriver struct {
	config     Config
	transferer transfer.ImageTransferer
	blobs      *blobs
	uploads    uploads
	manifests  *manifests
	metrics    tally.Scope
}

// NewReadWriteStorageDriver creates a KrakenStorageDriver which can push / pull blobs.
func NewReadWriteStorageDriver(
	config Config,
	cas *store.CAStore,
	transferer transfer.ImageTransferer,
	metrics tally.Scope) *KrakenStorageDriver {

	return &KrakenStorageDriver{
		config:     config,
		transferer: transferer,
		blobs:      newBlobs(cas, transferer),
		uploads:    newCASUploads(cas, transferer),
		manifests:  newManifests(transferer),
		metrics:    metrics,
	}
}

// NewReadOnlyStorageDriver creates a KrakenStorageDriver which can only pull blobs.
func NewReadOnlyStorageDriver(
	config Config,
	bs BlobStore,
	transferer transfer.ImageTransferer,
	metrics tally.Scope) *KrakenStorageDriver {

	return &KrakenStorageDriver{
		config:     config,
		transferer: transferer,
		blobs:      newBlobs(bs, transferer),
		uploads:    disabledUploads{},
		manifests:  newManifests(transferer),
		metrics:    metrics,
	}
}

// Name returns driver namae
func (d *KrakenStorageDriver) Name() string {
	return Name
}

// GetContent returns content in the path
// sample path: /docker/registry/v2/repositories/external/ubuntu/_layers/sha256/a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4/link
func (d *KrakenStorageDriver) GetContent(ctx context.Context, path string) (data []byte, err error) {
	log.Debugf("(*KrakenStorageDriver).GetContent %s", path)
	pathType, pathSubType, err := ParsePath(path)
	if err != nil {
		return nil, err
	}

	switch pathType {
	case _manifests:
		return d.manifests.getDigest(path, pathSubType)
	case _uploads:
		return d.uploads.getContent(path, pathSubType)
	case _layers:
		return d.blobs.getDigest(path)
	case _blobs:
		return d.blobs.getContent(ctx, path)
	}
	return nil, InvalidRequestError{path}
}

// Reader returns a reader of path at offset
func (d *KrakenStorageDriver) Reader(ctx context.Context, path string, offset int64) (reader io.ReadCloser, err error) {
	log.Debugf("(*KrakenStorageDriver).Reader %s", path)
	pathType, pathSubType, err := ParsePath(path)
	if err != nil {
		return nil, err
	}

	switch pathType {
	case _uploads:
		return d.uploads.reader(path, pathSubType, offset)
	case _blobs:
		return d.blobs.reader(ctx, path, offset)
	default:
		return nil, InvalidRequestError{path}
	}
}

// PutContent writes content to path
func (d *KrakenStorageDriver) PutContent(ctx context.Context, path string, content []byte) error {
	log.Debugf("(*KrakenStorageDriver).PutContent %s", path)
	pathType, pathSubType, err := ParsePath(path)
	if err != nil {
		return err
	}

	switch pathType {
	case _manifests:
		return d.manifests.putContent(path, pathSubType)
	case _uploads:
		return d.uploads.putContent(path, pathSubType, content)
	case _layers:
		// noop
		return nil
	case _blobs:
		return d.uploads.putBlobContent(path, content)
	default:
		return InvalidRequestError{path}
	}
}

// Writer returns a writer of path
func (d *KrakenStorageDriver) Writer(ctx context.Context, path string, append bool) (storagedriver.FileWriter, error) {
	log.Debugf("(*KrakenStorageDriver).Writer %s", path)
	pathType, pathSubType, err := ParsePath(path)
	if err != nil {
		return nil, err
	}

	switch pathType {
	case _uploads:
		w, err := d.uploads.writer(path, pathSubType)
		if err != nil {
			return nil, err
		}
		if append {
			if _, err := w.Seek(0, io.SeekEnd); err != nil {
				return nil, err
			}
		}
		return w, nil
	default:
		return nil, InvalidRequestError{path}
	}
}

// Stat returns fileinfo of path
func (d *KrakenStorageDriver) Stat(ctx context.Context, path string) (storagedriver.FileInfo, error) {
	log.Debugf("(*KrakenStorageDriver).Stat %s", path)
	pathType, _, err := ParsePath(path)
	if err != nil {
		return nil, err
	}

	switch pathType {
	case _uploads:
		return d.uploads.stat(path)
	case _blobs:
		return d.blobs.stat(ctx, path)
	case _manifests:
		return d.manifests.stat(path)
	default:
		return nil, InvalidRequestError{path}
	}
}

// List returns a list of content given path
func (d *KrakenStorageDriver) List(ctx context.Context, path string) ([]string, error) {
	log.Debugf("(*KrakenStorageDriver).List %s", path)
	pathType, pathSubType, err := ParsePath(path)
	if err != nil {
		return nil, err
	}

	switch pathType {
	case _uploads:
		return d.uploads.list(path, pathSubType)
	case _manifests:
		return d.manifests.list(path)
	default:
		return nil, InvalidRequestError{path}
	}
}

// Move moves sourcePath to destPath
func (d *KrakenStorageDriver) Move(ctx context.Context, sourcePath string, destPath string) error {
	log.Debugf("(*KrakenStorageDriver).Move %s %s", sourcePath, destPath)
	pathType, _, err := ParsePath(sourcePath)
	if err != nil {
		return err
	}

	switch pathType {
	case _uploads:
		return d.uploads.move(sourcePath, destPath)
	default:
		return InvalidRequestError{sourcePath + " to " + destPath}
	}
}

// Delete deletes path
func (d *KrakenStorageDriver) Delete(ctx context.Context, path string) error {
	log.Debugf("(*KrakenStorageDriver).Delete %s", path)
	return storagedriver.PathNotFoundError{
		DriverName: "p2p",
		Path:       path,
	}
}

// URLFor returns url for path
func (d *KrakenStorageDriver) URLFor(ctx context.Context, path string, options map[string]interface{}) (string, error) {
	log.Debugf("(*KrakenStorageDriver).URLFor %s", path)
	return "", fmt.Errorf("Not implemented")
}
<commit_msg>Remove unused consts from dockerregistry (#132)<commit_after>// Copyright (c) 2016-2019 Uber Technologies, 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 dockerregistry

import (
	"fmt"
	"io"

	"github.com/uber/kraken/lib/dockerregistry/transfer"
	"github.com/uber/kraken/lib/store"
	"github.com/uber/kraken/utils/log"

	"github.com/docker/distribution/context"
	storagedriver "github.com/docker/distribution/registry/storage/driver"
	"github.com/docker/distribution/registry/storage/driver/factory"
	"github.com/uber-go/tally"
)

// The path layout in the storage backend is roughly as follows:
//
//		<root>/v2
//			-> repositories/
// 				-><name>/
// 					-> _manifests/
// 						revisions
//							-> <manifest digest path>
//								-> link
// 						tags/<tag>
//							-> current/link
// 							-> index
//								-> <algorithm>/<hex digest>/link
// 					-> _layers/
// 						<layer links to blob store>
// 					-> _uploads/<id>
// 						data
// 						startedat
// 						hashstates/<algorithm>/<offset>
//			-> blobs/<algorithm>
//				<split directory content addressable storage>

// Name of storage driver.
const Name = "kraken"

func init() {
	factory.Register(Name, &krakenStorageDriverFactory{})
}

// InvalidRequestError implements error and contains the path that is not supported
type InvalidRequestError struct {
	path string
}

func (e InvalidRequestError) Error() string {
	return fmt.Sprintf("invalid request: %s", e.path)
}

type krakenStorageDriverFactory struct{}

func getParam(params map[string]interface{}, name string) interface{} {
	p, ok := params[name]
	if !ok || p == nil {
		log.Fatalf("Required parameter %s not found", name)
	}
	return p
}

func (factory *krakenStorageDriverFactory) Create(
	params map[string]interface{}) (storagedriver.StorageDriver, error) {

	// Common parameters.
	constructor := getParam(params, "constructor").(string)
	config := getParam(params, "config").(Config)
	transferer := getParam(params, "transferer").(transfer.ImageTransferer)
	metrics := getParam(params, "metrics").(tally.Scope)

	switch constructor {
	case _rw:
		castore := getParam(params, "castore").(*store.CAStore)
		return NewReadWriteStorageDriver(config, castore, transferer, metrics), nil
	case _ro:
		blobstore := getParam(params, "blobstore").(BlobStore)
		return NewReadOnlyStorageDriver(config, blobstore, transferer, metrics), nil
	default:
		return nil, fmt.Errorf("unknown constructor %s", constructor)
	}
}

// KrakenStorageDriver is a storage driver
type KrakenStorageDriver struct {
	config     Config
	transferer transfer.ImageTransferer
	blobs      *blobs
	uploads    uploads
	manifests  *manifests
	metrics    tally.Scope
}

// NewReadWriteStorageDriver creates a KrakenStorageDriver which can push / pull blobs.
func NewReadWriteStorageDriver(
	config Config,
	cas *store.CAStore,
	transferer transfer.ImageTransferer,
	metrics tally.Scope) *KrakenStorageDriver {

	return &KrakenStorageDriver{
		config:     config,
		transferer: transferer,
		blobs:      newBlobs(cas, transferer),
		uploads:    newCASUploads(cas, transferer),
		manifests:  newManifests(transferer),
		metrics:    metrics,
	}
}

// NewReadOnlyStorageDriver creates a KrakenStorageDriver which can only pull blobs.
func NewReadOnlyStorageDriver(
	config Config,
	bs BlobStore,
	transferer transfer.ImageTransferer,
	metrics tally.Scope) *KrakenStorageDriver {

	return &KrakenStorageDriver{
		config:     config,
		transferer: transferer,
		blobs:      newBlobs(bs, transferer),
		uploads:    disabledUploads{},
		manifests:  newManifests(transferer),
		metrics:    metrics,
	}
}

// Name returns driver namae
func (d *KrakenStorageDriver) Name() string {
	return Name
}

// GetContent returns content in the path
// sample path: /docker/registry/v2/repositories/external/ubuntu/_layers/sha256/a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4/link
func (d *KrakenStorageDriver) GetContent(ctx context.Context, path string) (data []byte, err error) {
	log.Debugf("(*KrakenStorageDriver).GetContent %s", path)
	pathType, pathSubType, err := ParsePath(path)
	if err != nil {
		return nil, err
	}

	switch pathType {
	case _manifests:
		return d.manifests.getDigest(path, pathSubType)
	case _uploads:
		return d.uploads.getContent(path, pathSubType)
	case _layers:
		return d.blobs.getDigest(path)
	case _blobs:
		return d.blobs.getContent(ctx, path)
	}
	return nil, InvalidRequestError{path}
}

// Reader returns a reader of path at offset
func (d *KrakenStorageDriver) Reader(ctx context.Context, path string, offset int64) (reader io.ReadCloser, err error) {
	log.Debugf("(*KrakenStorageDriver).Reader %s", path)
	pathType, pathSubType, err := ParsePath(path)
	if err != nil {
		return nil, err
	}

	switch pathType {
	case _uploads:
		return d.uploads.reader(path, pathSubType, offset)
	case _blobs:
		return d.blobs.reader(ctx, path, offset)
	default:
		return nil, InvalidRequestError{path}
	}
}

// PutContent writes content to path
func (d *KrakenStorageDriver) PutContent(ctx context.Context, path string, content []byte) error {
	log.Debugf("(*KrakenStorageDriver).PutContent %s", path)
	pathType, pathSubType, err := ParsePath(path)
	if err != nil {
		return err
	}

	switch pathType {
	case _manifests:
		return d.manifests.putContent(path, pathSubType)
	case _uploads:
		return d.uploads.putContent(path, pathSubType, content)
	case _layers:
		// noop
		return nil
	case _blobs:
		return d.uploads.putBlobContent(path, content)
	default:
		return InvalidRequestError{path}
	}
}

// Writer returns a writer of path
func (d *KrakenStorageDriver) Writer(ctx context.Context, path string, append bool) (storagedriver.FileWriter, error) {
	log.Debugf("(*KrakenStorageDriver).Writer %s", path)
	pathType, pathSubType, err := ParsePath(path)
	if err != nil {
		return nil, err
	}

	switch pathType {
	case _uploads:
		w, err := d.uploads.writer(path, pathSubType)
		if err != nil {
			return nil, err
		}
		if append {
			if _, err := w.Seek(0, io.SeekEnd); err != nil {
				return nil, err
			}
		}
		return w, nil
	default:
		return nil, InvalidRequestError{path}
	}
}

// Stat returns fileinfo of path
func (d *KrakenStorageDriver) Stat(ctx context.Context, path string) (storagedriver.FileInfo, error) {
	log.Debugf("(*KrakenStorageDriver).Stat %s", path)
	pathType, _, err := ParsePath(path)
	if err != nil {
		return nil, err
	}

	switch pathType {
	case _uploads:
		return d.uploads.stat(path)
	case _blobs:
		return d.blobs.stat(ctx, path)
	case _manifests:
		return d.manifests.stat(path)
	default:
		return nil, InvalidRequestError{path}
	}
}

// List returns a list of content given path
func (d *KrakenStorageDriver) List(ctx context.Context, path string) ([]string, error) {
	log.Debugf("(*KrakenStorageDriver).List %s", path)
	pathType, pathSubType, err := ParsePath(path)
	if err != nil {
		return nil, err
	}

	switch pathType {
	case _uploads:
		return d.uploads.list(path, pathSubType)
	case _manifests:
		return d.manifests.list(path)
	default:
		return nil, InvalidRequestError{path}
	}
}

// Move moves sourcePath to destPath
func (d *KrakenStorageDriver) Move(ctx context.Context, sourcePath string, destPath string) error {
	log.Debugf("(*KrakenStorageDriver).Move %s %s", sourcePath, destPath)
	pathType, _, err := ParsePath(sourcePath)
	if err != nil {
		return err
	}

	switch pathType {
	case _uploads:
		return d.uploads.move(sourcePath, destPath)
	default:
		return InvalidRequestError{sourcePath + " to " + destPath}
	}
}

// Delete deletes path
func (d *KrakenStorageDriver) Delete(ctx context.Context, path string) error {
	log.Debugf("(*KrakenStorageDriver).Delete %s", path)
	return storagedriver.PathNotFoundError{
		DriverName: "p2p",
		Path:       path,
	}
}

// URLFor returns url for path
func (d *KrakenStorageDriver) URLFor(ctx context.Context, path string, options map[string]interface{}) (string, error) {
	log.Debugf("(*KrakenStorageDriver).URLFor %s", path)
	return "", fmt.Errorf("Not implemented")
}
<|endoftext|>


<commit_before>package nessie

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestDoRequest(t *testing.T) {
	// Test structure to be serialized.
	type payload struct {
		A int `json:"a"`
	}
	authToken := "some token"
	var tests = []struct {
		method       string
		resource     string
		sentPayload  payload
		wantPayload  string
		serverStatus int
		wantStatus   []int
		wantError    bool
	}{
		// All succeeding methods.
		{"GET", "/test", payload{}, "{\"a\":0}", http.StatusOK, []int{http.StatusOK}, false},
		{"POST", "/test", payload{}, "{\"a\":0}", http.StatusOK, []int{http.StatusOK}, false},
		{"DELETE", "/test", payload{}, "{\"a\":0}", http.StatusOK, []int{http.StatusOK}, false},
		{"PUT", "/test", payload{}, "{\"a\":0}", http.StatusOK, []int{http.StatusOK}, false},
		// Payload test.
		{"GET", "/test", payload{42}, "{\"a\":42}", http.StatusOK, []int{http.StatusOK}, false},
		// Expected failure.
		{"POST", "/test", payload{}, "{\"a\":0}", http.StatusInternalServerError, []int{http.StatusInternalServerError}, false},
		// Unexpected failure
		{"POST", "/test", payload{}, "{\"a\":0}", http.StatusInternalServerError, []int{http.StatusOK}, true},
	}
	for _, tt := range tests {
		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(tt.serverStatus)
			if r.Header.Get("X-Cookie") != fmt.Sprintf("token=%s", authToken) {
				t.Errorf("invalid auth header, got=%s, want=%s", r.Header.Get("X-Cookie"), authToken)
			}
			body, err := ioutil.ReadAll(r.Body)
			if err != nil {
				t.Errorf("could not read request body: %v", err)
				return
			}
			bodyStr := string(body)
			if bodyStr != tt.wantPayload {
				t.Errorf("unexpected payload, got=%s, want=%s", body, tt.wantPayload)
			}
		}))
		n, err := NewInsecureNessus(ts.URL)
		n.Verbose = true
		if err != nil {
			t.Errorf("could not create nessie instance: %v (%+v)", err, tt)
			continue
		}
		// Increase covered lines.
		n.authCookie = authToken
		resp, err := n.doRequest(tt.method, tt.resource, tt.sentPayload, tt.wantStatus)
		if tt.wantError {
			if err == nil {
				t.Errorf("got no error, expected one (%+v)", tt)
			}
			continue
		}
		if err != nil {
			t.Errorf("error in doRequest: %v (%+v)", err, tt)
			continue
		}
		if resp.StatusCode != tt.serverStatus {
			t.Errorf("got status code=%d, wanted=%d", resp.StatusCode, tt.serverStatus)
		}
	}
}

func TestLogin(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(200)
		w.Header().Set("Content-Type", "application/json")
		j, err := json.Marshal(&loginResp{Token: "some token"})
		if err != nil {
			t.Fatalf("cannot serialize login response: %v", err)
		}
		w.Write(j)
	}))
	defer server.Close()
	n, err := NewInsecureNessus(server.URL)
	if err != nil {
		t.Fatalf("cannot create nessus instance: %v", err)
	}

	if err := n.Login("username", "password"); err != nil {
		t.Fatalf("got error during login: %v", err)
	}
	if got, want := n.authCookie, "some token"; got != want {
		t.Fatalf("wrong auth cookie, got=%q, want=%q", got, want)
	}
}

func TestMethods(t *testing.T) {
	var tests = []struct {
		resp       interface{}
		statusCode int
		call       func(n *Nessus)
	}{
		{&Session{}, http.StatusOK, func(n *Nessus) { n.Session() }},
		{&ServerProperties{}, http.StatusOK, func(n *Nessus) { n.ServerProperties() }},
		{&ServerStatus{}, http.StatusOK, func(n *Nessus) { n.ServerStatus() }},
		{&User{}, http.StatusOK, func(n *Nessus) {
			n.CreateUser("username", "pass", UserTypeLocal, Permissions32, "name", "email@foo.com")
		}},
		{&listUsersResp{}, http.StatusOK, func(n *Nessus) { n.ListUsers() }},
		{nil, http.StatusOK, func(n *Nessus) { n.DeleteUser(42) }},
		{nil, http.StatusOK, func(n *Nessus) { n.SetUserPassword(42, "newpass") }},
		{&User{}, http.StatusOK, func(n *Nessus) {
			n.EditUser(42, Permissions128, "newname", "newmain@goo.fom")
		}},
		{[]PluginFamily{}, http.StatusOK, func(n *Nessus) { n.PluginFamilies() }},
	}
	for _, tt := range tests {
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(tt.statusCode)
			if tt.resp != nil {
				j, err := json.Marshal(tt.resp)
				if err != nil {
					t.Fatalf("cannot serialize response: %v", err)
				}
				w.Write(j)
			}
		}))
		defer server.Close()
		n, err := NewInsecureNessus(server.URL)
		if err != nil {
			t.Fatalf("cannot create nessus instance: %v", err)
		}
		n.Verbose = true
		tt.call(n)
	}
}
<commit_msg>second batch of tests for all methods<commit_after>package nessie

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestDoRequest(t *testing.T) {
	// Test structure to be serialized.
	type payload struct {
		A int `json:"a"`
	}
	authToken := "some token"
	var tests = []struct {
		method       string
		resource     string
		sentPayload  payload
		wantPayload  string
		serverStatus int
		wantStatus   []int
		wantError    bool
	}{
		// All succeeding methods.
		{"GET", "/test", payload{}, "{\"a\":0}", http.StatusOK, []int{http.StatusOK}, false},
		{"POST", "/test", payload{}, "{\"a\":0}", http.StatusOK, []int{http.StatusOK}, false},
		{"DELETE", "/test", payload{}, "{\"a\":0}", http.StatusOK, []int{http.StatusOK}, false},
		{"PUT", "/test", payload{}, "{\"a\":0}", http.StatusOK, []int{http.StatusOK}, false},
		// Payload test.
		{"GET", "/test", payload{42}, "{\"a\":42}", http.StatusOK, []int{http.StatusOK}, false},
		// Expected failure.
		{"POST", "/test", payload{}, "{\"a\":0}", http.StatusInternalServerError, []int{http.StatusInternalServerError}, false},
		// Unexpected failure
		{"POST", "/test", payload{}, "{\"a\":0}", http.StatusInternalServerError, []int{http.StatusOK}, true},
	}
	for _, tt := range tests {
		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(tt.serverStatus)
			if r.Header.Get("X-Cookie") != fmt.Sprintf("token=%s", authToken) {
				t.Errorf("invalid auth header, got=%s, want=%s", r.Header.Get("X-Cookie"), authToken)
			}
			body, err := ioutil.ReadAll(r.Body)
			if err != nil {
				t.Errorf("could not read request body: %v", err)
				return
			}
			bodyStr := string(body)
			if bodyStr != tt.wantPayload {
				t.Errorf("unexpected payload, got=%s, want=%s", body, tt.wantPayload)
			}
		}))
		n, err := NewInsecureNessus(ts.URL)
		n.Verbose = true
		if err != nil {
			t.Errorf("could not create nessie instance: %v (%+v)", err, tt)
			continue
		}
		// Increase covered lines.
		n.authCookie = authToken
		resp, err := n.doRequest(tt.method, tt.resource, tt.sentPayload, tt.wantStatus)
		if tt.wantError {
			if err == nil {
				t.Errorf("got no error, expected one (%+v)", tt)
			}
			continue
		}
		if err != nil {
			t.Errorf("error in doRequest: %v (%+v)", err, tt)
			continue
		}
		if resp.StatusCode != tt.serverStatus {
			t.Errorf("got status code=%d, wanted=%d", resp.StatusCode, tt.serverStatus)
		}
	}
}

func TestLogin(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(200)
		w.Header().Set("Content-Type", "application/json")
		j, err := json.Marshal(&loginResp{Token: "some token"})
		if err != nil {
			t.Fatalf("cannot serialize login response: %v", err)
		}
		w.Write(j)
	}))
	defer server.Close()
	n, err := NewInsecureNessus(server.URL)
	if err != nil {
		t.Fatalf("cannot create nessus instance: %v", err)
	}

	if err := n.Login("username", "password"); err != nil {
		t.Fatalf("got error during login: %v", err)
	}
	if got, want := n.authCookie, "some token"; got != want {
		t.Fatalf("wrong auth cookie, got=%q, want=%q", got, want)
	}
}

func TestMethods(t *testing.T) {
	var tests = []struct {
		resp       interface{}
		statusCode int
		call       func(n *Nessus)
	}{
		{&Session{}, http.StatusOK, func(n *Nessus) { n.Session() }},
		{&ServerProperties{}, http.StatusOK, func(n *Nessus) { n.ServerProperties() }},
		{&ServerStatus{}, http.StatusOK, func(n *Nessus) { n.ServerStatus() }},
		{&User{}, http.StatusOK, func(n *Nessus) {
			n.CreateUser("username", "pass", UserTypeLocal, Permissions32, "name", "email@foo.com")
		}},
		{&listUsersResp{}, http.StatusOK, func(n *Nessus) { n.ListUsers() }},
		{nil, http.StatusOK, func(n *Nessus) { n.DeleteUser(42) }},
		{nil, http.StatusOK, func(n *Nessus) { n.SetUserPassword(42, "newpass") }},
		{&User{}, http.StatusOK, func(n *Nessus) {
			n.EditUser(42, Permissions128, "newname", "newmain@goo.fom")
		}},
		{[]PluginFamily{}, http.StatusOK, func(n *Nessus) { n.PluginFamilies() }},
		{&FamilyDetails{}, http.StatusOK, func(n *Nessus) { n.FamilyDetails(42) }},
		{&PluginDetails{}, http.StatusOK, func(n *Nessus) { n.PluginDetails(42) }},
		{[]Scanner{}, http.StatusOK, func(n *Nessus) { n.Scanners() }},
		{&listPoliciesResp{}, http.StatusOK, func(n *Nessus) { n.Policies() }},
		{&Scan{}, http.StatusOK, func(n *Nessus) {
			n.NewScan("editorUUID", "settingsName", 42, 43, 44, LaunchDaily, []string{"target1", "target2"})
		}},
		{&ListScansResponse{}, http.StatusOK, func(n *Nessus) { n.Scans() }},
		{[]Template{}, http.StatusOK, func(n *Nessus) { n.ScanTemplates() }},
		{[]Template{}, http.StatusOK, func(n *Nessus) { n.PolicyTemplates() }},
		{"id", http.StatusOK, func(n *Nessus) { n.StartScan(42) }},
		{nil, http.StatusOK, func(n *Nessus) { n.PauseScan(42) }},
		{nil, http.StatusOK, func(n *Nessus) { n.ResumeScan(42) }},
		{nil, http.StatusOK, func(n *Nessus) { n.StopScan(42) }},
		{&ScanDetailsResp{}, http.StatusOK, func(n *Nessus) { n.ScanDetails(42) }},
		{[]TimeZone{}, http.StatusOK, func(n *Nessus) { n.Timezones() }},
		{[]Folder{}, http.StatusOK, func(n *Nessus) { n.Folders() }},
		{nil, http.StatusOK, func(n *Nessus) { n.CreateFolder("name") }},
		{nil, http.StatusOK, func(n *Nessus) { n.EditFolder(42, "newname") }},
		{nil, http.StatusOK, func(n *Nessus) { n.DeleteFolder(42) }},
		{42, http.StatusOK, func(n *Nessus) { n.ExportScan(42, ExportPDF) }},
		{true, http.StatusOK, func(n *Nessus) { n.ExportFinished(42, 43) }},
		{[]byte("raw export"), http.StatusOK, func(n *Nessus) { n.DownloadExport(42, 43) }},
		{[]Permission{}, http.StatusOK, func(n *Nessus) { n.Permissions("scanner", 42) }},
	}
	for _, tt := range tests {
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(tt.statusCode)
			if tt.resp != nil {
				j, err := json.Marshal(tt.resp)
				if err != nil {
					t.Fatalf("cannot serialize response: %v", err)
				}
				w.Write(j)
			}
		}))
		defer server.Close()
		n, err := NewInsecureNessus(server.URL)
		if err != nil {
			t.Fatalf("cannot create nessus instance: %v", err)
		}
		n.Verbose = true
		tt.call(n)
	}
}
<|endoftext|>


<commit_before>package oauth

import (
	"net/url"
	"testing"
)

const (
	baseServer = "http://oauthbin.appspot.com"
)

var (
	requestToken = baseServer + "/v1/request-token"
	accessToken  = baseServer + "/v1/access-token"
	echo         = baseServer + "/v1/echo"
)

// These tests use the testing oAuth server
// documented at http://term.ie/oauth/example/

func testOAuth(t *testing.T, method string, values url.Values) {
	c := &Consumer{
		Key:             "key",
		Secret:          "secret",
		RequestTokenURL: requestToken,
		AccessTokenURL:  accessToken,
		CallbackURL:     "oob",
	}
	_, rt, err := c.Authorization()
	if err != nil {
		t.Fatal(err)
	}
	t.Logf("request token is %+v", rt)
	at, err := c.Exchange(rt, "")
	if err != nil {
		t.Fatal(err)
	}
	t.Logf("access token is %+v", at)
	t.Logf("sending request with values %+v", values)
	resp, err := c.SendRequest(method, echo, values, at)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Close()
	data, err := resp.ReadAll()
	if err != nil {
		t.Fatal(err)
	}
	s := string(data)
	t.Logf("server replied with %q", s)
	if len(values) == 0 {
		if s != "" {
			t.Errorf("expected empty response, got %q", s)
		}
	} else {
		if e := values.Encode(); e != s {
			t.Errorf("expected %q got %q", e, s)
		}
	}
}

func TestOAuth(t *testing.T) {
	testOAuth(t, "GET", nil)
}

func TestGet(t *testing.T) {
	testOAuth(t, "GET", url.Values{"foo": []string{"bar"}})
}

func TestPost(t *testing.T) {
	testOAuth(t, "POST", url.Values{"foo": []string{"bar"}})
}

func TestUnicode(t *testing.T) {
	testOAuth(t, "POST", url.Values{"alberto": []string{"garcía"}})
}
<commit_msg>Add an oAuth test for problematic characters<commit_after>package oauth

import (
	"net/url"
	"reflect"
	"testing"
)

const (
	baseServer = "http://oauthbin.appspot.com"
)

var (
	requestToken = baseServer + "/v1/request-token"
	accessToken  = baseServer + "/v1/access-token"
	echo         = baseServer + "/v1/echo"
)

// These tests use the testing oAuth server
// documented at http://term.ie/oauth/example/

func testOAuth(t *testing.T, method string, values url.Values) {
	c := &Consumer{
		Key:             "key",
		Secret:          "secret",
		RequestTokenURL: requestToken,
		AccessTokenURL:  accessToken,
		CallbackURL:     "oob",
	}
	_, rt, err := c.Authorization()
	if err != nil {
		t.Fatal(err)
	}
	t.Logf("request token is %+v", rt)
	at, err := c.Exchange(rt, "")
	if err != nil {
		t.Fatal(err)
	}
	t.Logf("access token is %+v", at)
	t.Logf("sending request with values %+v", values)
	resp, err := c.SendRequest(method, echo, values, at)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Close()
	data, err := resp.ReadAll()
	if err != nil {
		t.Fatal(err)
	}
	s := string(data)
	t.Logf("server replied with %q", s)
	if len(values) == 0 {
		if s != "" {
			t.Errorf("expected empty response, got %q", s)
		}
	} else {
		rec, err := url.ParseQuery(s)
		if err != nil {
			t.Error("error parsing received values: %s", err)
		}
		if !reflect.DeepEqual(values, rec) {
			t.Errorf("expecting values %v, got %v instead", values, rec)
		}
	}
}

func TestOAuth(t *testing.T) {
	testOAuth(t, "GET", nil)
}

func TestGet(t *testing.T) {
	testOAuth(t, "GET", url.Values{"foo": []string{"bar"}})
}

func TestPost(t *testing.T) {
	testOAuth(t, "POST", url.Values{"foo": []string{"bar"}})
}

func TestUnicode(t *testing.T) {
	testOAuth(t, "POST", url.Values{"alberto": []string{"garcía"}})
}

func TestProblematicChars(t *testing.T) {
	values := make(url.Values)
	values.Add("a", "=")
	values.Add("b", "+/*")
	values.Add("c", "~")
	testOAuth(t, "GET", values)
}
<|endoftext|>


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

package rocserv

import (
	"fmt"
	"time"
	"encoding/json"
	"strconv"
	"sort"
	"crypto/sha1"
	"crypto/md5"

	// now use 73a8ef737e8ea002281a28b4cb92a1de121ad4c6
    "github.com/coreos/go-etcd/etcd"

	"github.com/sdming/gosnow"

	"github.com/shawnfeng/sutil"
	"github.com/shawnfeng/sutil/slog"

	"github.com/shawnfeng/roc/util/dbrouter"
)

// id生成器
// 服务注册

type ServInfo struct {
	Type string         `json:"type"`
	Addr string         `json:"addr"`
	//Processor string    `json:"processor"`
}

func (m *ServInfo) String() string {
	return fmt.Sprintf("type:%s addr:%s", m.Type, m.Addr)
}

// ServBase Interface
type ServBase interface {
	// key is processor to ServInfo
	RegisterService(servs map[string]*ServInfo) error
	Servid() int
	// 服务副本名称, servename + servid
	Copyname() string

	// 获取服务的配置
	ServConfig(cfg interface{}) error
	// 任意路径的配置信息
	//ArbiConfig(location string) (string, error)

	// id生成逻辑
	GenSnowFlakeId() (uint64, error)
	GenUuid() string
	GenUuidSha1() string
	GenUuidMd5() string

	// db router
	Dbrouter() *dbrouter.Router
}

//====================
// id生成逻辑
type IdGenerator struct {
	snow *gosnow.SnowFlake
}

func (m *IdGenerator) GenSnowFlakeId() (uint64, error) {
	return m.snow.Next()
}

func (m *IdGenerator) GenUuid() string {
	return sutil.GetUUID()
}

func (m *IdGenerator) GenUuidSha1() string {
	h := sha1.Sum([]byte(m.GenUuid()))
	return fmt.Sprintf("%x", h)
}

func (m *IdGenerator) GenUuidMd5() string {
	h := md5.Sum([]byte(m.GenUuid()))
	return fmt.Sprintf("%x", h)
}

//====================================
func getValue(client *etcd.Client, path string) ([]byte, error) {
    r, err := client.Get(path, false, false)
	if err != nil {
		return nil, err
	}

	if r.Node == nil || r.Node.Dir {
		return nil, fmt.Errorf("etcd node value err location:%s", path)
	}

	return []byte(r.Node.Value), nil

}

func genSid(client *etcd.Client, path, skey string) (int, error) {
	fun := "genSid -->"
    r, err := client.Get(path, false, false)
	if err != nil {
		return -1, err
	}

	js, _ := json.Marshal(r)

	slog.Infof("%s", js)

	if r.Node == nil || !r.Node.Dir {
		return -1, fmt.Errorf("node error location:%s", path)
	}

	slog.Infof("%s serv:%s len:%d", fun, r.Node.Key, r.Node.Nodes.Len())

	// 获取已有的servid，按从小到大排列
	ids := make([]int, 0)
	for _, n := range r.Node.Nodes {
		sid := n.Key[len(r.Node.Key)+1:]
		id, err := strconv.Atoi(sid)
		if err != nil || id < 0 {
			slog.Errorf("%s sid error key:%s", fun, n.Key)
		} else {
			ids = append(ids, id)
			if n.Value == skey {
				// 如果已经存在的sid使用的skey和设置一致，则使用之前的sid
				return id, nil
			}
		}
	}

	sort.Ints(ids)
	sid := 0
	for _, id := range ids {
		// 取不重复的最小的id
		if sid == id {
			sid++
		} else {
			break
		}
	}

	nserv := fmt.Sprintf("%s/%d", r.Node.Key, sid)
	r, err = client.Create(nserv, skey, 0)
	if err != nil {
		return -1, err
	}

	jr, _ := json.Marshal(r)
	slog.Infof("%s newserv:%s rep:%s", fun, nserv, jr)

	return sid, nil

}

func retryGenSid(client *etcd.Client, path, skey string, try int) (int, error) {
	fun := "retryGenSid -->"
	for i := 0; i < try; i++ {
		// 重试3次
		sid, err := genSid(client, path, skey)
		if err != nil {
			slog.Errorf("%s gensid try:%d path:%s err:%s", fun, i, path, err)
		} else {
			return sid, nil
		}
	}

	return -1, fmt.Errorf("gensid error try:%d", try)
}

func initSnowflake(servid int) (*gosnow.SnowFlake, error) {
	if servid < 0 {
		return nil, fmt.Errorf("init snowflake use nagtive servid")
	}
	gosnow.Since = time.Date(2014, 11, 1, 0, 0, 0, 0, time.UTC).UnixNano() / 1000000
	v, err := gosnow.NewSnowFlake(uint32(servid))
	if err != nil {
		return nil, err
	}

	return v, nil
}

<commit_msg>snowflake id use int64<commit_after>// Copyright 2014 The roc Author. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rocserv

import (
	"fmt"
	"time"
	"encoding/json"
	"strconv"
	"sort"
	"crypto/sha1"
	"crypto/md5"

	// now use 73a8ef737e8ea002281a28b4cb92a1de121ad4c6
    "github.com/coreos/go-etcd/etcd"

	"github.com/sdming/gosnow"

	"github.com/shawnfeng/sutil"
	"github.com/shawnfeng/sutil/slog"

	"github.com/shawnfeng/roc/util/dbrouter"
)

// id生成器
// 服务注册

type ServInfo struct {
	Type string         `json:"type"`
	Addr string         `json:"addr"`
	//Processor string    `json:"processor"`
}

func (m *ServInfo) String() string {
	return fmt.Sprintf("type:%s addr:%s", m.Type, m.Addr)
}

// ServBase Interface
type ServBase interface {
	// key is processor to ServInfo
	RegisterService(servs map[string]*ServInfo) error
	Servid() int
	// 服务副本名称, servename + servid
	Copyname() string

	// 获取服务的配置
	ServConfig(cfg interface{}) error
	// 任意路径的配置信息
	//ArbiConfig(location string) (string, error)

	// id生成逻辑
	GenSnowFlakeId() (int64, error)
	GenUuid() string
	GenUuidSha1() string
	GenUuidMd5() string

	// db router
	Dbrouter() *dbrouter.Router
}

//====================
// id生成逻辑
type IdGenerator struct {
	snow *gosnow.SnowFlake
}

func (m *IdGenerator) GenSnowFlakeId() (int64, error) {
	id, err := m.snow.Next()
	return int64(id), err
}

func (m *IdGenerator) GenUuid() string {
	return sutil.GetUUID()
}

func (m *IdGenerator) GenUuidSha1() string {
	h := sha1.Sum([]byte(m.GenUuid()))
	return fmt.Sprintf("%x", h)
}

func (m *IdGenerator) GenUuidMd5() string {
	h := md5.Sum([]byte(m.GenUuid()))
	return fmt.Sprintf("%x", h)
}

//====================================
func getValue(client *etcd.Client, path string) ([]byte, error) {
    r, err := client.Get(path, false, false)
	if err != nil {
		return nil, err
	}

	if r.Node == nil || r.Node.Dir {
		return nil, fmt.Errorf("etcd node value err location:%s", path)
	}

	return []byte(r.Node.Value), nil

}

func genSid(client *etcd.Client, path, skey string) (int, error) {
	fun := "genSid -->"
    r, err := client.Get(path, false, false)
	if err != nil {
		return -1, err
	}

	js, _ := json.Marshal(r)

	slog.Infof("%s", js)

	if r.Node == nil || !r.Node.Dir {
		return -1, fmt.Errorf("node error location:%s", path)
	}

	slog.Infof("%s serv:%s len:%d", fun, r.Node.Key, r.Node.Nodes.Len())

	// 获取已有的servid，按从小到大排列
	ids := make([]int, 0)
	for _, n := range r.Node.Nodes {
		sid := n.Key[len(r.Node.Key)+1:]
		id, err := strconv.Atoi(sid)
		if err != nil || id < 0 {
			slog.Errorf("%s sid error key:%s", fun, n.Key)
		} else {
			ids = append(ids, id)
			if n.Value == skey {
				// 如果已经存在的sid使用的skey和设置一致，则使用之前的sid
				return id, nil
			}
		}
	}

	sort.Ints(ids)
	sid := 0
	for _, id := range ids {
		// 取不重复的最小的id
		if sid == id {
			sid++
		} else {
			break
		}
	}

	nserv := fmt.Sprintf("%s/%d", r.Node.Key, sid)
	r, err = client.Create(nserv, skey, 0)
	if err != nil {
		return -1, err
	}

	jr, _ := json.Marshal(r)
	slog.Infof("%s newserv:%s rep:%s", fun, nserv, jr)

	return sid, nil

}

func retryGenSid(client *etcd.Client, path, skey string, try int) (int, error) {
	fun := "retryGenSid -->"
	for i := 0; i < try; i++ {
		// 重试3次
		sid, err := genSid(client, path, skey)
		if err != nil {
			slog.Errorf("%s gensid try:%d path:%s err:%s", fun, i, path, err)
		} else {
			return sid, nil
		}
	}

	return -1, fmt.Errorf("gensid error try:%d", try)
}

func initSnowflake(servid int) (*gosnow.SnowFlake, error) {
	if servid < 0 {
		return nil, fmt.Errorf("init snowflake use nagtive servid")
	}
	gosnow.Since = time.Date(2014, 11, 1, 0, 0, 0, 0, time.UTC).UnixNano() / 1000000
	v, err := gosnow.NewSnowFlake(uint32(servid))
	if err != nil {
		return nil, err
	}

	return v, nil
}

<|endoftext|>


<commit_before>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published
   by the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

package manager

import (
	"io"
	"mime/multipart"
	"path"
	"path/filepath"
	"strings"

	uploadClient "github.com/webx-top/client/upload"
	_ "github.com/webx-top/client/upload/driver"
	"github.com/webx-top/echo"

	"github.com/admpub/nging/application/handler"
	"github.com/admpub/nging/application/library/common"
	modelFile "github.com/admpub/nging/application/model/file"
	"github.com/admpub/nging/application/registry/upload"
	"github.com/admpub/nging/application/registry/upload/driver/filesystem"
	"github.com/admpub/nging/application/registry/upload/helper"
	"github.com/admpub/nging/application/registry/upload/convert"
	"github.com/admpub/qrcode"
)

// 文件上传保存路径规则：
// 表名称/表行ID/文件名

// ResponseDataForUpload 根据不同的上传方式响应不同的数据格式
func ResponseDataForUpload(ctx echo.Context, field string, err error, imageURLs []string) (result echo.H, embed bool) {
	return upload.ResponserGet(field)(ctx, field, err, imageURLs)
}

var (
	StorerEngine = filesystem.Name
)

func File(ctx echo.Context) error {
	uploadType := ctx.Param(`type`)
	typ, _, _ := getTableInfo(uploadType)
	file := ctx.Param(`*`)
	file = filepath.Join(helper.UploadDir, typ, file)
	extension := filepath.Ext(file)
	extension = strings.ToLower(extension)
	convert, ok := convert.GetConverter(extension)
	if !ok {
		return ctx.File(file)
	}
	var supported bool
	for _, accept := range ctx.Accept().Type {
		if accept.Mime == `image/webp` {
			supported = true
			break
		}
	}
	originalFile := strings.TrimSuffix(file, extension)
	if !supported {
		return ctx.File(originalFile)
	}
	newStore := upload.StorerGet(StorerEngine)
	if newStore == nil {
		return ctx.E(`存储引擎“%s”未被登记`, StorerEngine)
	}
	f, err := os.Open(originalFile)
	if err != nil {
		return echo.ErrNotFound
	}
	defer f.Close()
	buf, err := convert(f, 70)
	if err != nil {
		return err
	}
	storer := newStore(ctx, typ)
	var savePath, viewURL string
	savePath, viewURL, err = storer.Put(storer.URLToFile(file), buf, int64(len(buf.Bytes())))
	_, _ = savePath, viewURL
	return ctx.ServeContent(f, path.Base(file), time.Now())
}

// SaveFilename SaveFilename(`0/`,``,`img.jpg`)
func SaveFilename(subdir, name, postFilename string) (string, error) {
	ext := filepath.Ext(postFilename)
	fname := name
	if len(fname) == 0 {
		var err error
		fname, err = common.UniqueID()
		if err != nil {
			return ``, err
		}
	}
	fname += ext
	return subdir + fname, nil
}

// Upload 上传文件
func Upload(ctx echo.Context) error {
	ownerType := `user`
	user := handler.User(ctx)
	var ownerID uint64
	if user != nil {
		ownerID = uint64(user.Id)
	}
	if ownerID < 1 {
		ctx.Data().SetError(ctx.E(`请先登录`))
		return ctx.Redirect(handler.URLFor(`/login`))
	}
	return UploadByOwner(ctx, ownerType, ownerID)
}

// UploadByOwner 上传文件
func UploadByOwner(ctx echo.Context, ownerType string, ownerID uint64) error {
	uploadType := ctx.Param(`type`)
	field := ctx.Query(`field`) // 上传表单file输入框名称
	pipe := ctx.Form(`pipe`)
	var (
		err      error
		fileURLs []string
	)
	if len(uploadType) == 0 {
		err = ctx.E(`请提供参数“%s”`, ctx.Path())
		datax, embed := ResponseDataForUpload(ctx, field, err, fileURLs)
		if !embed {
			return ctx.JSON(datax)
		}
		return err
	}
	fileType := ctx.Form(`filetype`)
	prepareData, err := upload.Prepare(ctx, uploadType, fileType, StorerEngine)
	if err != nil {
		datax, embed := ResponseDataForUpload(ctx, field, err, fileURLs)
		if !embed {
			return ctx.JSON(datax)
		}
	}
	storer := prepareData.Storer(ctx)
	defer prepareData.Close()
	fileM := modelFile.NewFile(ctx)
	fileM.StorerName = StorerEngine
	fileM.TableId = ``
	fileM.SetFieldName(prepareData.FieldName)
	fileM.SetTableName(prepareData.TableName)
	fileM.OwnerId = ownerID
	fileM.OwnerType = ownerType
	fileM.Type = fileType

	subdir, name, err := prepareData.Checkin(ctx, fileM)
	if err != nil {
		datax, embed := ResponseDataForUpload(ctx, field, err, fileURLs)
		if !embed {
			return ctx.JSON(datax)
		}
		return err
	}

	clientName := ctx.Form(`client`)
	if len(clientName) > 0 {
		result := &uploadClient.Result{}
		result.SetFileNameGenerator(func(filename string) (string, error) {
			return SaveFilename(subdir, name, filename)
		})

		client := uploadClient.Upload(ctx, clientName, result, storer, watermarkFile, prepareData.Checker)
		if client.GetError() != nil {
			if client.GetError() == upload.ErrExistsFile {
				client.SetError(nil)
			}
			return client.Response()
		}

		fileM.SetByUploadResult(result)

		var reader io.ReadCloser
		reader, err = storer.Get(result.SavePath)
		if reader != nil {
			defer reader.Close()
		}
		if err != nil {
			return client.SetError(err).Response()
		}
		err = prepareData.DBSaver(fileM, result, reader)
		return client.SetError(err).Response()
	}
	var results uploadClient.Results
	results, err = upload.BatchUpload(
		ctx,
		`files[]`,
		func(r *uploadClient.Result) (string, error) {
			if err := prepareData.Checker(r); err != nil {
				return ``, err
			}
			return SaveFilename(subdir, name, r.FileName)
		},
		storer,
		func(result *uploadClient.Result, file multipart.File) error {
			fileM.Id = 0
			fileM.SetByUploadResult(result)
			return prepareData.DBSaver(fileM, result, file)
		},
		watermarkFile,
	)
	datax, embed := ResponseDataForUpload(ctx, field, err, results.FileURLs())
	if err != nil {
		if !embed {
			return ctx.JSON(datax)
		}
		return err
	}

	if pipe == `deqr` { //解析二维码
		if len(results) > 0 {
			reader, err := storer.Get(results[0].SavePath)
			if reader != nil {
				defer reader.Close()
			}
			if err != nil {
				if !embed {
					datax[`raw`] = err.Error()
					return ctx.JSON(datax)
				}
				return err
			}
			raw, err := qrcode.Decode(reader, strings.TrimPrefix(path.Ext(results[0].SavePath), `.`))
			if err != nil {
				raw = err.Error()
			}
			datax[`raw`] = raw
		}
	}
	if !embed {
		return ctx.JSON(datax)
	}
	data := ctx.Data()
	data.SetData(datax)
	return ctx.JSON(data)
}

func getTableInfo(uploadType string) (tableName string, fieldName string, defaults []string) {
	return upload.GetTableInfo(uploadType)
}
<commit_msg>update<commit_after>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published
   by the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

package manager

import (
	"os"
	"time"
	"io"
	"mime/multipart"
	"path"
	"path/filepath"
	"strings"

	uploadClient "github.com/webx-top/client/upload"
	_ "github.com/webx-top/client/upload/driver"
	"github.com/webx-top/echo"

	"github.com/admpub/nging/application/handler"
	"github.com/admpub/nging/application/library/common"
	modelFile "github.com/admpub/nging/application/model/file"
	"github.com/admpub/nging/application/registry/upload"
	"github.com/admpub/nging/application/registry/upload/driver/filesystem"
	"github.com/admpub/nging/application/registry/upload/helper"
	"github.com/admpub/nging/application/registry/upload/convert"
	"github.com/admpub/qrcode"
)

// 文件上传保存路径规则：
// 表名称/表行ID/文件名

// ResponseDataForUpload 根据不同的上传方式响应不同的数据格式
func ResponseDataForUpload(ctx echo.Context, field string, err error, imageURLs []string) (result echo.H, embed bool) {
	return upload.ResponserGet(field)(ctx, field, err, imageURLs)
}

var (
	StorerEngine = filesystem.Name
)

func File(ctx echo.Context) error {
	uploadType := ctx.Param(`type`)
	typ, _, _ := getTableInfo(uploadType)
	file := ctx.Param(`*`)
	file = filepath.Join(helper.UploadDir, typ, file)
	extension := filepath.Ext(file)
	extension = strings.ToLower(extension)
	convert, ok := convert.GetConverter(extension)
	if !ok {
		return ctx.File(file)
	}
	var supported bool
	for _, accept := range ctx.Accept().Type {
		if accept.Mime == `image/webp` {
			supported = true
			break
		}
	}
	originalFile := strings.TrimSuffix(file, extension)
	if !supported {
		return ctx.File(originalFile)
	}
	newStore := upload.StorerGet(StorerEngine)
	if newStore == nil {
		return ctx.E(`存储引擎“%s”未被登记`, StorerEngine)
	}
	f, err := os.Open(originalFile)
	if err != nil {
		return echo.ErrNotFound
	}
	defer f.Close()
	buf, err := convert(f, 70)
	if err != nil {
		return err
	}
	storer := newStore(ctx, typ)
	var savePath, viewURL string
	savePath, viewURL, err = storer.Put(storer.URLToFile(file), buf, int64(len(buf.Bytes())))
	_, _ = savePath, viewURL
	return ctx.ServeContent(f, path.Base(file), time.Now())
}

// SaveFilename SaveFilename(`0/`,``,`img.jpg`)
func SaveFilename(subdir, name, postFilename string) (string, error) {
	ext := filepath.Ext(postFilename)
	fname := name
	if len(fname) == 0 {
		var err error
		fname, err = common.UniqueID()
		if err != nil {
			return ``, err
		}
	}
	fname += ext
	return subdir + fname, nil
}

// Upload 上传文件
func Upload(ctx echo.Context) error {
	ownerType := `user`
	user := handler.User(ctx)
	var ownerID uint64
	if user != nil {
		ownerID = uint64(user.Id)
	}
	if ownerID < 1 {
		ctx.Data().SetError(ctx.E(`请先登录`))
		return ctx.Redirect(handler.URLFor(`/login`))
	}
	return UploadByOwner(ctx, ownerType, ownerID)
}

// UploadByOwner 上传文件
func UploadByOwner(ctx echo.Context, ownerType string, ownerID uint64) error {
	uploadType := ctx.Param(`type`)
	field := ctx.Query(`field`) // 上传表单file输入框名称
	pipe := ctx.Form(`pipe`)
	var (
		err      error
		fileURLs []string
	)
	if len(uploadType) == 0 {
		err = ctx.E(`请提供参数“%s”`, ctx.Path())
		datax, embed := ResponseDataForUpload(ctx, field, err, fileURLs)
		if !embed {
			return ctx.JSON(datax)
		}
		return err
	}
	fileType := ctx.Form(`filetype`)
	prepareData, err := upload.Prepare(ctx, uploadType, fileType, StorerEngine)
	if err != nil {
		datax, embed := ResponseDataForUpload(ctx, field, err, fileURLs)
		if !embed {
			return ctx.JSON(datax)
		}
	}
	storer := prepareData.Storer(ctx)
	defer prepareData.Close()
	fileM := modelFile.NewFile(ctx)
	fileM.StorerName = StorerEngine
	fileM.TableId = ``
	fileM.SetFieldName(prepareData.FieldName)
	fileM.SetTableName(prepareData.TableName)
	fileM.OwnerId = ownerID
	fileM.OwnerType = ownerType
	fileM.Type = fileType

	subdir, name, err := prepareData.Checkin(ctx, fileM)
	if err != nil {
		datax, embed := ResponseDataForUpload(ctx, field, err, fileURLs)
		if !embed {
			return ctx.JSON(datax)
		}
		return err
	}

	clientName := ctx.Form(`client`)
	if len(clientName) > 0 {
		result := &uploadClient.Result{}
		result.SetFileNameGenerator(func(filename string) (string, error) {
			return SaveFilename(subdir, name, filename)
		})

		client := uploadClient.Upload(ctx, clientName, result, storer, watermarkFile, prepareData.Checker)
		if client.GetError() != nil {
			if client.GetError() == upload.ErrExistsFile {
				client.SetError(nil)
			}
			return client.Response()
		}

		fileM.SetByUploadResult(result)

		var reader io.ReadCloser
		reader, err = storer.Get(result.SavePath)
		if reader != nil {
			defer reader.Close()
		}
		if err != nil {
			return client.SetError(err).Response()
		}
		err = prepareData.DBSaver(fileM, result, reader)
		return client.SetError(err).Response()
	}
	var results uploadClient.Results
	results, err = upload.BatchUpload(
		ctx,
		`files[]`,
		func(r *uploadClient.Result) (string, error) {
			if err := prepareData.Checker(r); err != nil {
				return ``, err
			}
			return SaveFilename(subdir, name, r.FileName)
		},
		storer,
		func(result *uploadClient.Result, file multipart.File) error {
			fileM.Id = 0
			fileM.SetByUploadResult(result)
			return prepareData.DBSaver(fileM, result, file)
		},
		watermarkFile,
	)
	datax, embed := ResponseDataForUpload(ctx, field, err, results.FileURLs())
	if err != nil {
		if !embed {
			return ctx.JSON(datax)
		}
		return err
	}

	if pipe == `deqr` { //解析二维码
		if len(results) > 0 {
			reader, err := storer.Get(results[0].SavePath)
			if reader != nil {
				defer reader.Close()
			}
			if err != nil {
				if !embed {
					datax[`raw`] = err.Error()
					return ctx.JSON(datax)
				}
				return err
			}
			raw, err := qrcode.Decode(reader, strings.TrimPrefix(path.Ext(results[0].SavePath), `.`))
			if err != nil {
				raw = err.Error()
			}
			datax[`raw`] = raw
		}
	}
	if !embed {
		return ctx.JSON(datax)
	}
	data := ctx.Data()
	data.SetData(datax)
	return ctx.JSON(data)
}

func getTableInfo(uploadType string) (tableName string, fieldName string, defaults []string) {
	return upload.GetTableInfo(uploadType)
}
<|endoftext|>


<commit_before>// +build linux,!nonsystemd

package journald

import (
	"context"
	"sync"
	"time"

	"github.com/coreos/go-systemd/sdjournal"
	"github.com/inconshreveable/log15"
)

var Supported bool = true

type JournaldReader interface {
	Start()
	Stop()
	Entries() chan map[string]string
}

type reader struct {
	journal  *sdjournal.Journal
	entries  chan map[string]string
	stopchan chan bool
	wgroup   *sync.WaitGroup
	logger   log15.Logger
}

func NewReader(ctx context.Context, logger log15.Logger) (JournaldReader, error) {
	var err error
	r := &reader{logger: logger}
	r.journal, err = sdjournal.NewJournal()
	if err != nil {
		return nil, err
	}
	err = r.journal.SeekTail()
	if err != nil {
		r.journal.Close()
		return nil, err
	}
	_, err = r.journal.Previous()
	if err != nil {
		r.journal.Close()
		return nil, err
	}
	r.entries = make(chan map[string]string)
	r.wgroup = &sync.WaitGroup{}

	go func() {
		<-ctx.Done()
		r.Stop()
		close(r.entries)
		r.journal.Close()
	}()

	return r, nil
}

func (r *reader) Entries() chan map[string]string {
	return r.entries
}

func (r *reader) wait() chan int {
	events := make(chan int)
	r.wgroup.Add(1)

	go func() {
		defer r.wgroup.Done()
		var ev int

	WaitLoop:
		for {
			select {
			case <-r.stopchan:
				break WaitLoop
			default:
				ev = r.journal.Wait(time.Second)
				if ev == sdjournal.SD_JOURNAL_APPEND || ev == sdjournal.SD_JOURNAL_INVALIDATE {
					events <- ev
					close(events)
					break WaitLoop
				}
			}
		}
	}()

	return events
}

func (r *reader) Start() {
	r.stopchan = make(chan bool)
	r.wgroup.Add(1)

	go func() {
		defer r.wgroup.Done()
		var err error
		var nb uint64
		var entry *sdjournal.JournalEntry
		for {
			// get entries from journald
		LoopGetEntries:
			for {
				select {
				case <-r.stopchan:
					return
				default:
					nb, err = r.journal.Next()
					if err != nil {
						r.logger.Warn("journal.Next() error", "error", err)
					} else if nb == 0 {
						r.logger.Debug("0 entry in the journal")
						break LoopGetEntries
					} else {
						entry, err = r.journal.GetEntry()
						if err != nil {
							r.logger.Warn("journal.GetEntry() error", "error", err)
						} else {
							r.entries <- entry.Fields
						}
					}
				}
			}

			// wait that journald has more entries
			r.logger.Debug("Waiting for more journal entries")
			events := r.wait()
			select {
			case <-events:
			case <-r.stopchan:
				return
			}
			r.logger.Debug("There are more journal entries")
		}
	}()
}

func (r *reader) Stop() {
	if r.stopchan != nil {
		close(r.stopchan)
	}
	r.wgroup.Wait()
}
<commit_msg>less debug logs<commit_after>// +build linux,!nonsystemd

package journald

import (
	"context"
	"sync"
	"time"

	"github.com/coreos/go-systemd/sdjournal"
	"github.com/inconshreveable/log15"
)

var Supported bool = true

type JournaldReader interface {
	Start()
	Stop()
	Entries() chan map[string]string
}

type reader struct {
	journal  *sdjournal.Journal
	entries  chan map[string]string
	stopchan chan bool
	wgroup   *sync.WaitGroup
	logger   log15.Logger
}

func NewReader(ctx context.Context, logger log15.Logger) (JournaldReader, error) {
	var err error
	r := &reader{logger: logger}
	r.journal, err = sdjournal.NewJournal()
	if err != nil {
		return nil, err
	}
	err = r.journal.SeekTail()
	if err != nil {
		r.journal.Close()
		return nil, err
	}
	_, err = r.journal.Previous()
	if err != nil {
		r.journal.Close()
		return nil, err
	}
	r.entries = make(chan map[string]string)
	r.wgroup = &sync.WaitGroup{}

	go func() {
		<-ctx.Done()
		r.Stop()
		close(r.entries)
		r.journal.Close()
	}()

	return r, nil
}

func (r *reader) Entries() chan map[string]string {
	return r.entries
}

func (r *reader) wait() chan int {
	events := make(chan int)
	r.wgroup.Add(1)

	go func() {
		defer r.wgroup.Done()
		var ev int

	WaitLoop:
		for {
			select {
			case <-r.stopchan:
				break WaitLoop
			default:
				ev = r.journal.Wait(time.Second)
				if ev == sdjournal.SD_JOURNAL_APPEND || ev == sdjournal.SD_JOURNAL_INVALIDATE {
					events <- ev
					close(events)
					break WaitLoop
				}
			}
		}
	}()

	return events
}

func (r *reader) Start() {
	r.stopchan = make(chan bool)
	r.wgroup.Add(1)

	go func() {
		defer r.wgroup.Done()
		var err error
		var nb uint64
		var entry *sdjournal.JournalEntry
		for {
			// get entries from journald
		LoopGetEntries:
			for {
				select {
				case <-r.stopchan:
					return
				default:
					nb, err = r.journal.Next()
					if err != nil {
						r.logger.Warn("journal.Next() error", "error", err)
					} else if nb == 0 {
						break LoopGetEntries
					} else {
						entry, err = r.journal.GetEntry()
						if err != nil {
							r.logger.Warn("journal.GetEntry() error", "error", err)
						} else {
							r.entries <- entry.Fields
						}
					}
				}
			}

			// wait that journald has more entries
			events := r.wait()
			select {
			case <-events:
			case <-r.stopchan:
				return
			}
		}
	}()
}

func (r *reader) Stop() {
	if r.stopchan != nil {
		close(r.stopchan)
	}
	r.wgroup.Wait()
}
<|endoftext|>


<commit_before>package ParquetFile

import (
	"context"

	"cloud.google.com/go/storage"
)

type GcsFile struct {
	ProjectId  string
	BucketName string
	Ctx        context.Context

	Client     *storage.Client
	Bucket     *storage.BucketHandle
	FilePath   string
	FileReader *storage.Reader
	FileWriter *storage.Writer
}

func NewGcsFileWriter(projectId string, bucketName string, name string, ctx context.Context) (ParquetFile, error) {
	res := &GcsFile{
		ProjectId:  projectId,
		BucketName: bucketName,
		Ctx:        ctx,
		FilePath:   name,
	}
	return res.Create(name)
}

func NewGcsFileReader(hosts []string, user string, name string) (ParquetFile, error) {
	res := &GcsFile{
		FilePath: name,
	}
	return res.Open(name)
}

func (self *GcsFile) Create(name string) (ParquetFile, error) {
	var err error
	gcs := new(GcsFile)
	gcs.Client, err = storage.NewClient(self.Ctx)
	gcs.FilePath = name
	if err != nil {
		return gcs, err
	}
	// must use existing bucket
	gcs.Bucket = gcs.Client.Bucket(self.BucketName)
	obj := gcs.Bucket.Object(name)
	gcs.FileWriter = obj.NewWriter(self.Ctx)
	return gcs, err

}
func (self *GcsFile) Open(name string) (ParquetFile, error) {
	var err error
	gcs := new(GcsFile)
	gcs.Client, err = storage.NewClient(self.Ctx)
	gcs.FilePath = name
	if err != nil {
		return gcs, err
	}
	// must use existing bucket
	gcs.Bucket = gcs.Client.Bucket(self.BucketName)
	obj := gcs.Bucket.Object(name)
	gcs.FileReader, err = obj.NewReader(self.Ctx)
	return gcs, err
}
func (self *GcsFile) Seek(offset int, pos int) (int64, error) {
	//Not implemented
	return 0, nil
}

func (self *GcsFile) Read(b []byte) (n int, err error) {
	return self.FileReader.Read(b)
}

func (self *GcsFile) Write(b []byte) (n int, err error) {
	return self.FileWriter.Write(b)
}

func (self *GcsFile) Close() {
	if self.FileReader != nil {
		self.FileReader.Close()
	}
	if self.FileWriter != nil {
		self.FileWriter.Close()
	}
	if self.Client != nil {
		self.Client.Close()
	}
}
<commit_msg>Update gcs close function implementation<commit_after>package ParquetFile

import (
	"context"

	"cloud.google.com/go/storage"
)

type GcsFile struct {
	ProjectId  string
	BucketName string
	Ctx        context.Context

	Client     *storage.Client
	Bucket     *storage.BucketHandle
	FilePath   string
	FileReader *storage.Reader
	FileWriter *storage.Writer
}

func NewGcsFileWriter(projectId string, bucketName string, name string, ctx context.Context) (ParquetFile, error) {
	res := &GcsFile{
		ProjectId:  projectId,
		BucketName: bucketName,
		Ctx:        ctx,
		FilePath:   name,
	}
	return res.Create(name)
}

func NewGcsFileReader(projectId string, bucketName string, name string, ctx context.Context) (ParquetFile, error) {
	res := &GcsFile{
		ProjectId:  projectId,
		BucketName: bucketName,
		Ctx:        ctx,
		FilePath:   name,
	}
	return res.Open(name)
}

func (self *GcsFile) Create(name string) (ParquetFile, error) {
	var err error
	gcs := new(GcsFile)
	gcs.Client, err = storage.NewClient(self.Ctx)
	gcs.FilePath = name
	if err != nil {
		return gcs, err
	}
	// must use existing bucket
	gcs.Bucket = gcs.Client.Bucket(self.BucketName)
	obj := gcs.Bucket.Object(name)
	gcs.FileWriter = obj.NewWriter(self.Ctx)
	return gcs, err

}
func (self *GcsFile) Open(name string) (ParquetFile, error) {
	var err error
	gcs := new(GcsFile)
	gcs.Client, err = storage.NewClient(self.Ctx)
	gcs.FilePath = name
	if err != nil {
		return gcs, err
	}
	// must use existing bucket
	gcs.Bucket = gcs.Client.Bucket(self.BucketName)
	obj := gcs.Bucket.Object(name)
	gcs.FileReader, err = obj.NewReader(self.Ctx)
	return gcs, err
}
func (self *GcsFile) Seek(offset int64, pos int) (int64, error) {
	//Not implemented
	return 0, nil
}

func (self *GcsFile) Read(b []byte) (n int, err error) {
	return self.FileReader.Read(b)
}

func (self *GcsFile) Write(b []byte) (n int, err error) {
	return self.FileWriter.Write(b)
}

func (self *GcsFile) Close() error {
	if self.FileReader != nil {
		if err := self.FileReader.Close(); err != nil {
			return err
		}
	}
	if self.FileWriter != nil {
		if err := self.FileWriter.Close(); err != nil {
			return err
		}
	}
	if self.Client != nil {
		if err := self.Client.Close(); err != nil {
			return nil
		}
	}
	return nil
}
<|endoftext|>