<commit_before>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"strings"
	"sync/atomic"
	"time"
	"unsafe"

	r "github.com/dancannon/gorethink"
	"github.com/jaracil/ei"
	. "github.com/jaracil/nexus/log"
	"github.com/jaracil/nxcli/nxcore"
	"github.com/jaracil/smartio"
	"golang.org/x/net/context"
)

type JsonRpcErr struct {
	Code    int         `json:"code"`
	Message string      `json:"message"`
	Data    interface{} `json:"data,omitempty"`
}

type JsonRpcReq struct {
	Jsonrpc string      `json:"jsonrpc"`
	Id      interface{} `json:"id,omitempty"`
	Method  string      `json:"method"`
	Params  interface{} `json:"params"`
	nc      *NexusConn
}

type JsonRpcRes struct {
	Jsonrpc string      `json:"jsonrpc"`
	Id      interface{} `json:"id,omitempty"`
	Result  interface{} `json:"result,omitempty"`
	Error   *JsonRpcErr `json:"error,omitempty"`
}

type NexusConn struct {
	conn      net.Conn
	proto     string
	connRx    *smartio.SmartReader
	connTx    *smartio.SmartWriter
	connId    string
	user      *UserData
	chRes     chan *JsonRpcRes
	chReq     chan *JsonRpcReq
	context   context.Context
	cancelFun context.CancelFunc
	wdog      int64
	closed    int32 //Atomic bool
}

func NewNexusConn(conn net.Conn) *NexusConn {
	nc := &NexusConn{
		conn:   conn,
		proto:  "unknown",
		connRx: smartio.NewSmartReader(conn),
		connTx: smartio.NewSmartWriter(conn),
		connId: nodeId + safeId(4),
		user:   Nobody,
		chRes:  make(chan *JsonRpcRes, 64),
		chReq:  make(chan *JsonRpcReq, 64),
		wdog:   90,
	}
	nc.context, nc.cancelFun = context.WithCancel(mainContext)
	return nc
}

func NewInternalClient() *nxcore.NexusConn {
	client, server := net.Pipe()

	nc := NewNexusConn(server)
	nc.proto = "internal"
	nc.user = &UserData{
		User: "@internal",
		Tags: map[string]map[string]interface{}{
			".": map[string]interface{}{
				"@admin": true,
			},
		},
	}
	go nc.handle()

	return nxcore.NewNexusConn(client)
}

func (req *JsonRpcReq) Error(code int, message string, data interface{}) {
	if code < 0 {
		if message != "" {
			message = fmt.Sprintf("%s:[%s]", ErrStr[code], message)
		} else {
			message = ErrStr[code]
		}
	}
	req.nc.pushRes(
		&JsonRpcRes{
			Id:    req.Id,
			Error: &JsonRpcErr{Code: code, Message: message, Data: data},
		},
	)
}

func (req *JsonRpcReq) Result(result interface{}) {
	req.nc.pushRes(
		&JsonRpcRes{
			Id:     req.Id,
			Result: result,
		},
	)
}

func (nc *NexusConn) pushRes(res *JsonRpcRes) (err error) {
	select {
	case nc.chRes <- res:
	case <-nc.context.Done():
		err = errors.New("Context cancelled")
	}
	return
}

func (nc *NexusConn) pullRes() (res *JsonRpcRes, err error) {
	select {
	case res = <-nc.chRes:
	case <-nc.context.Done():
		err = errors.New("Context cancelled")
	}
	return
}

func (nc *NexusConn) pushReq(req *JsonRpcReq) (err error) {
	select {
	case nc.chReq <- req:
	case <-nc.context.Done():
		err = errors.New("Context cancelled")
	}
	return
}

func (nc *NexusConn) pullReq() (req *JsonRpcReq, err error) {
	select {
	case req = <-nc.chReq:
	case <-nc.context.Done():
		err = errors.New("Context cancelled")
	}
	return
}

func getTags(ud *UserData, prefix string) (tags map[string]interface{}) {
	tags = map[string]interface{}{}
	if ud == nil || ud.Tags == nil {
		return
	}
	pfs := prefixes(prefix)
	for _, pf := range pfs {
		if tm, ok := ud.Tags[pf]; ok {
			for k, v := range tm {
				tags[k] = v
			}
		}
	}
	return
}

func (nc *NexusConn) getTags(prefix string) (tags map[string]interface{}) {
	return getTags(nc.user, prefix)
}

func (nc *NexusConn) handleReq(req *JsonRpcReq) {
	if req.Id == nil {
		return
	}
	switch {
	case strings.HasPrefix(req.Method, "sys."):
		switch {
		case strings.HasPrefix(req.Method, "sys.node."):
			nc.handleNodesReq(req)
		case strings.HasPrefix(req.Method, "sys.session."):
			nc.handleSessionReq(req)
		default:
			nc.handleSysReq(req)
		}
	case strings.HasPrefix(req.Method, "task."):
		nc.handleTaskReq(req)
	case strings.HasPrefix(req.Method, "pipe."):
		nc.handlePipeReq(req)
	case strings.HasPrefix(req.Method, "topic."):
		nc.handleTopicReq(req)
	case strings.HasPrefix(req.Method, "user."):
		nc.handleUserReq(req)
	case strings.HasPrefix(req.Method, "sync."):
		nc.handleSyncReq(req)

	default:
		req.Error(ErrMethodNotFound, "", nil)
	}
}

func (nc *NexusConn) respWorker() {
	defer nc.close()
	trackCh, err := sesNotify.Register(nc.connId, make(chan interface{}, 1024))
	if err != nil { // Duplicated session ???
		Log.Warnf("Error on [%s] respWorker: %s", nc.connId, err)
		return
	}
	defer sesNotify.Unregister(nc.connId)
	for {
		select {
		case d := <-trackCh:

			switch res := d.(type) {

			case *Task:
				if res.ErrCode != nil {
					nc.pushRes(
						&JsonRpcRes{
							Id:    res.LocalId,
							Error: &JsonRpcErr{Code: *res.ErrCode, Message: res.ErrStr, Data: res.ErrObj},
						},
					)
				} else {
					nc.pushRes(
						&JsonRpcRes{
							Id:     res.LocalId,
							Result: res.Result,
						},
					)
				}

			case *Session:
				if res.Reload {
					nc.reload(false)
				}
				if res.Kick {
					Log.Printf("Connection [%s] has been kicked!", nc.connId)
					nc.close()
				}
			}

		case <-nc.context.Done():
			return
		}
	}
}

func (nc *NexusConn) sendWorker() {
	defer nc.close()
	var null *int
	for {
		res, err := nc.pullRes()
		if err != nil {
			Log.Debugf("Error on [%s] sendWorker: %s", nc.connId, err)
			break
		}
		if res.Id == nil {
			if res.Error == nil {
				continue //Skip notification responses
			}
			if res.Error.Code == ErrInvalidRequest || res.Error.Code == ErrParse {
				res.Id = null
			} else {
				continue
			}
		}
		res.Jsonrpc = "2.0"
		if res.Result == nil && res.Error == nil {
			res.Result = null
		}
		buf, err := json.Marshal(res)
		if err != nil {
			Log.Debugf("[%s] connection marshal error: %s", nc.connId, err)
			break
		}
		buf = append(buf, byte('\r'), byte('\n'))
		n, err := nc.connTx.Write(buf)
		if err != nil || n != len(buf) {
			Log.Debugf("[%s] connection write error: %s", nc.connId, err)
			break
		}
	}
	Log.Debugf("Exit from [%s] sendWorker", nc.connId)
}

func (nc *NexusConn) recvWorker() {
	defer nc.close()
	dec := json.NewDecoder(nc.connRx)
	for dec.More() {
		req := &JsonRpcReq{nc: nc}
		nc.connRx.SetLimit(1024 * 1024 * 32)
		err := dec.Decode(req)
		if err != nil {
			if _, ok := err.(*json.SyntaxError); ok {
				req.Error(ErrParse, "", nil)
				dec = json.NewDecoder(nc.connRx) // Refresh decoder
				continue
			}
			if _, ok := err.(*json.UnmarshalTypeError); ok {
				req.Error(ErrInvalidRequest, "", nil)
				continue
			}
			break
		}
		err = nc.pushReq(req)
		if err != nil {
			Log.Debugf("Error on [%s] recvWorker: %s", nc.connId, err)
			break
		}
	}
	Log.Debugf("Exit from [%s] recvWorker", nc.connId)
}

func (nc *NexusConn) watchdog() {
	defer nc.close()
	tick := time.NewTicker(time.Second * 10)
	exit := false
	for !exit {
		select {
		case <-tick.C:
			now := time.Now().Unix()
			max := atomic.LoadInt64(&nc.wdog)
			if (now-nc.connRx.GetLast() > max) &&
				(now-nc.connTx.GetLast() > max) {
				exit = true
				Log.Warnf("Connection [%s] watch dog expired!", nc.connId)
			}

			nc.updateSession()

		case <-nc.context.Done():
			exit = true
		}
	}
	tick.Stop()
}

func (nc *NexusConn) close() {
	if atomic.CompareAndSwapInt32(&nc.closed, 0, 1) {
		nc.cancelFun()
		nc.conn.Close()
		if mainContext.Err() == nil {
			Log.Printf("Closing [%s] session", nc.connId)
			dbClean(nc.connId)
		}
	}
}

func (nc *NexusConn) reload(fromSameSession bool) (bool, int) {
	if nc.user == nil {
		return false, ErrInvalidRequest
	}
	ud, err := loadUserData(nc.user.User)
	if err != ErrNoError {
		return false, ErrInternal
	}

	atomic.StorePointer((*unsafe.Pointer)(unsafe.Pointer(&nc.user)), unsafe.Pointer(&UserData{User: ud.User, Mask: ud.Mask, Tags: maskTags(ud.Tags, ud.Mask)}))

	if !fromSameSession {
		wres, err := r.Table("sessions").
			Between(nc.connId, nc.connId+"\uffff").
			Update(ei.M{"reload": false}).
			RunWrite(db)
		if err != nil || wres.Replaced == 0 {
			return false, ErrInternal
		}
		Log.Printf("Connection [%s] reloaded by other session", nc.connId)
	} else {
		Log.Printf("Connection [%s] reloaded by itself", nc.connId)
	}
	return true, 0
}

func (nc *NexusConn) updateSession() {
	_, err := r.Table("sessions").
		Get(nc.connId).
		Replace(ei.M{
			"id":            nc.connId,
			"nodeId":        nodeId,
			"creationTime":  r.Row.Field("creationTime").Default(r.Now()),
			"lastSeen":      r.Now(),
			"remoteAddress": nc.conn.RemoteAddr().String(),
			"protocol":      nc.proto,
			"user":          nc.user.User,
		}).
		RunWrite(db)

	if err != nil {
		Log.Errorf("Error updating session [%s]: %s", nc.connId, err)
		nc.close()
	}
}

var numconn int64

func (nc *NexusConn) handle() {

	if nc.proto != "internal" {
		atomic.AddInt64(&numconn, 1)
		defer func() { atomic.AddInt64(&numconn, -1) }()
	}

	defer nc.close()
	go nc.respWorker()
	go nc.sendWorker()
	go nc.recvWorker()
	go nc.watchdog()

	nc.updateSession()

	for {
		req, err := nc.pullReq()
		if err != nil {
			Log.Debugf("Error on [%s] connection handler: %s", nc.connId, err)
			break
		}
		Log.Printf("[%s@%s] %s: %#v - id: %.0f", req.nc.connId, req.nc.conn.RemoteAddr(), req.Method, req.Params, req.Id)
		if (req.Jsonrpc != "2.0" && req.Jsonrpc != "") || req.Method == "" { //"jsonrpc":"2.0" is optional
			req.Error(ErrInvalidRequest, "", nil)
			continue
		}
		go nc.handleReq(req)
	}
}
<commit_msg>Marshal the request parameters in the log message<commit_after>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"strings"
	"sync/atomic"
	"time"
	"unsafe"

	r "github.com/dancannon/gorethink"
	"github.com/jaracil/ei"
	. "github.com/jaracil/nexus/log"
	"github.com/jaracil/nxcli/nxcore"
	"github.com/jaracil/smartio"
	"golang.org/x/net/context"
)

type JsonRpcErr struct {
	Code    int         `json:"code"`
	Message string      `json:"message"`
	Data    interface{} `json:"data,omitempty"`
}

type JsonRpcReq struct {
	Jsonrpc string      `json:"jsonrpc"`
	Id      interface{} `json:"id,omitempty"`
	Method  string      `json:"method"`
	Params  interface{} `json:"params"`
	nc      *NexusConn
}

type JsonRpcRes struct {
	Jsonrpc string      `json:"jsonrpc"`
	Id      interface{} `json:"id,omitempty"`
	Result  interface{} `json:"result,omitempty"`
	Error   *JsonRpcErr `json:"error,omitempty"`
}

type NexusConn struct {
	conn      net.Conn
	proto     string
	connRx    *smartio.SmartReader
	connTx    *smartio.SmartWriter
	connId    string
	user      *UserData
	chRes     chan *JsonRpcRes
	chReq     chan *JsonRpcReq
	context   context.Context
	cancelFun context.CancelFunc
	wdog      int64
	closed    int32 //Atomic bool
}

func NewNexusConn(conn net.Conn) *NexusConn {
	nc := &NexusConn{
		conn:   conn,
		proto:  "unknown",
		connRx: smartio.NewSmartReader(conn),
		connTx: smartio.NewSmartWriter(conn),
		connId: nodeId + safeId(4),
		user:   Nobody,
		chRes:  make(chan *JsonRpcRes, 64),
		chReq:  make(chan *JsonRpcReq, 64),
		wdog:   90,
	}
	nc.context, nc.cancelFun = context.WithCancel(mainContext)
	return nc
}

func NewInternalClient() *nxcore.NexusConn {
	client, server := net.Pipe()

	nc := NewNexusConn(server)
	nc.proto = "internal"
	nc.user = &UserData{
		User: "@internal",
		Tags: map[string]map[string]interface{}{
			".": map[string]interface{}{
				"@admin": true,
			},
		},
	}
	go nc.handle()

	return nxcore.NewNexusConn(client)
}

func (req *JsonRpcReq) Error(code int, message string, data interface{}) {
	if code < 0 {
		if message != "" {
			message = fmt.Sprintf("%s:[%s]", ErrStr[code], message)
		} else {
			message = ErrStr[code]
		}
	}
	req.nc.pushRes(
		&JsonRpcRes{
			Id:    req.Id,
			Error: &JsonRpcErr{Code: code, Message: message, Data: data},
		},
	)
}

func (req *JsonRpcReq) Result(result interface{}) {
	req.nc.pushRes(
		&JsonRpcRes{
			Id:     req.Id,
			Result: result,
		},
	)
}

func (nc *NexusConn) pushRes(res *JsonRpcRes) (err error) {
	select {
	case nc.chRes <- res:
	case <-nc.context.Done():
		err = errors.New("Context cancelled")
	}
	return
}

func (nc *NexusConn) pullRes() (res *JsonRpcRes, err error) {
	select {
	case res = <-nc.chRes:
	case <-nc.context.Done():
		err = errors.New("Context cancelled")
	}
	return
}

func (nc *NexusConn) pushReq(req *JsonRpcReq) (err error) {
	select {
	case nc.chReq <- req:
	case <-nc.context.Done():
		err = errors.New("Context cancelled")
	}
	return
}

func (nc *NexusConn) pullReq() (req *JsonRpcReq, err error) {
	select {
	case req = <-nc.chReq:
	case <-nc.context.Done():
		err = errors.New("Context cancelled")
	}
	return
}

func getTags(ud *UserData, prefix string) (tags map[string]interface{}) {
	tags = map[string]interface{}{}
	if ud == nil || ud.Tags == nil {
		return
	}
	pfs := prefixes(prefix)
	for _, pf := range pfs {
		if tm, ok := ud.Tags[pf]; ok {
			for k, v := range tm {
				tags[k] = v
			}
		}
	}
	return
}

func (nc *NexusConn) getTags(prefix string) (tags map[string]interface{}) {
	return getTags(nc.user, prefix)
}

func (nc *NexusConn) handleReq(req *JsonRpcReq) {
	if req.Id == nil {
		return
	}
	switch {
	case strings.HasPrefix(req.Method, "sys."):
		switch {
		case strings.HasPrefix(req.Method, "sys.node."):
			nc.handleNodesReq(req)
		case strings.HasPrefix(req.Method, "sys.session."):
			nc.handleSessionReq(req)
		default:
			nc.handleSysReq(req)
		}
	case strings.HasPrefix(req.Method, "task."):
		nc.handleTaskReq(req)
	case strings.HasPrefix(req.Method, "pipe."):
		nc.handlePipeReq(req)
	case strings.HasPrefix(req.Method, "topic."):
		nc.handleTopicReq(req)
	case strings.HasPrefix(req.Method, "user."):
		nc.handleUserReq(req)
	case strings.HasPrefix(req.Method, "sync."):
		nc.handleSyncReq(req)

	default:
		req.Error(ErrMethodNotFound, "", nil)
	}
}

func (nc *NexusConn) respWorker() {
	defer nc.close()
	trackCh, err := sesNotify.Register(nc.connId, make(chan interface{}, 1024))
	if err != nil { // Duplicated session ???
		Log.Warnf("Error on [%s] respWorker: %s", nc.connId, err)
		return
	}
	defer sesNotify.Unregister(nc.connId)
	for {
		select {
		case d := <-trackCh:

			switch res := d.(type) {

			case *Task:
				if res.ErrCode != nil {
					nc.pushRes(
						&JsonRpcRes{
							Id:    res.LocalId,
							Error: &JsonRpcErr{Code: *res.ErrCode, Message: res.ErrStr, Data: res.ErrObj},
						},
					)
				} else {
					nc.pushRes(
						&JsonRpcRes{
							Id:     res.LocalId,
							Result: res.Result,
						},
					)
				}

			case *Session:
				if res.Reload {
					nc.reload(false)
				}
				if res.Kick {
					Log.Printf("Connection [%s] has been kicked!", nc.connId)
					nc.close()
				}
			}

		case <-nc.context.Done():
			return
		}
	}
}

func (nc *NexusConn) sendWorker() {
	defer nc.close()
	var null *int
	for {
		res, err := nc.pullRes()
		if err != nil {
			Log.Debugf("Error on [%s] sendWorker: %s", nc.connId, err)
			break
		}
		if res.Id == nil {
			if res.Error == nil {
				continue //Skip notification responses
			}
			if res.Error.Code == ErrInvalidRequest || res.Error.Code == ErrParse {
				res.Id = null
			} else {
				continue
			}
		}
		res.Jsonrpc = "2.0"
		if res.Result == nil && res.Error == nil {
			res.Result = null
		}
		buf, err := json.Marshal(res)
		if err != nil {
			Log.Debugf("[%s] connection marshal error: %s", nc.connId, err)
			break
		}
		buf = append(buf, byte('\r'), byte('\n'))
		n, err := nc.connTx.Write(buf)
		if err != nil || n != len(buf) {
			Log.Debugf("[%s] connection write error: %s", nc.connId, err)
			break
		}
	}
	Log.Debugf("Exit from [%s] sendWorker", nc.connId)
}

func (nc *NexusConn) recvWorker() {
	defer nc.close()
	dec := json.NewDecoder(nc.connRx)
	for dec.More() {
		req := &JsonRpcReq{nc: nc}
		nc.connRx.SetLimit(1024 * 1024 * 32)
		err := dec.Decode(req)
		if err != nil {
			if _, ok := err.(*json.SyntaxError); ok {
				req.Error(ErrParse, "", nil)
				dec = json.NewDecoder(nc.connRx) // Refresh decoder
				continue
			}
			if _, ok := err.(*json.UnmarshalTypeError); ok {
				req.Error(ErrInvalidRequest, "", nil)
				continue
			}
			break
		}
		err = nc.pushReq(req)
		if err != nil {
			Log.Debugf("Error on [%s] recvWorker: %s", nc.connId, err)
			break
		}
	}
	Log.Debugf("Exit from [%s] recvWorker", nc.connId)
}

func (nc *NexusConn) watchdog() {
	defer nc.close()
	tick := time.NewTicker(time.Second * 10)
	exit := false
	for !exit {
		select {
		case <-tick.C:
			now := time.Now().Unix()
			max := atomic.LoadInt64(&nc.wdog)
			if (now-nc.connRx.GetLast() > max) &&
				(now-nc.connTx.GetLast() > max) {
				exit = true
				Log.Warnf("Connection [%s] watch dog expired!", nc.connId)
			}

			nc.updateSession()

		case <-nc.context.Done():
			exit = true
		}
	}
	tick.Stop()
}

func (nc *NexusConn) close() {
	if atomic.CompareAndSwapInt32(&nc.closed, 0, 1) {
		nc.cancelFun()
		nc.conn.Close()
		if mainContext.Err() == nil {
			Log.Printf("Closing [%s] session", nc.connId)
			dbClean(nc.connId)
		}
	}
}

func (nc *NexusConn) reload(fromSameSession bool) (bool, int) {
	if nc.user == nil {
		return false, ErrInvalidRequest
	}
	ud, err := loadUserData(nc.user.User)
	if err != ErrNoError {
		return false, ErrInternal
	}

	atomic.StorePointer((*unsafe.Pointer)(unsafe.Pointer(&nc.user)), unsafe.Pointer(&UserData{User: ud.User, Mask: ud.Mask, Tags: maskTags(ud.Tags, ud.Mask)}))

	if !fromSameSession {
		wres, err := r.Table("sessions").
			Between(nc.connId, nc.connId+"\uffff").
			Update(ei.M{"reload": false}).
			RunWrite(db)
		if err != nil || wres.Replaced == 0 {
			return false, ErrInternal
		}
		Log.Printf("Connection [%s] reloaded by other session", nc.connId)
	} else {
		Log.Printf("Connection [%s] reloaded by itself", nc.connId)
	}
	return true, 0
}

func (nc *NexusConn) updateSession() {
	_, err := r.Table("sessions").
		Get(nc.connId).
		Replace(ei.M{
			"id":            nc.connId,
			"nodeId":        nodeId,
			"creationTime":  r.Row.Field("creationTime").Default(r.Now()),
			"lastSeen":      r.Now(),
			"remoteAddress": nc.conn.RemoteAddr().String(),
			"protocol":      nc.proto,
			"user":          nc.user.User,
		}).
		RunWrite(db)

	if err != nil {
		Log.Errorf("Error updating session [%s]: %s", nc.connId, err)
		nc.close()
	}
}

var numconn int64

func (nc *NexusConn) handle() {

	if nc.proto != "internal" {
		atomic.AddInt64(&numconn, 1)
		defer func() { atomic.AddInt64(&numconn, -1) }()
	}

	defer nc.close()
	go nc.respWorker()
	go nc.sendWorker()
	go nc.recvWorker()
	go nc.watchdog()

	nc.updateSession()

	for {
		req, err := nc.pullReq()
		if err != nil {
			Log.Debugf("Error on [%s] connection handler: %s", nc.connId, err)
			break
		}
		params, err := json.Marshal(req.Params)
		if err != nil {
			Log.Printf("[%s@%s] %s: %#v - id: %.0f", req.nc.connId, req.nc.conn.RemoteAddr(), req.Method, req.Params, req.Id)
		} else {
			Log.Printf("[%s@%s] %s: %s - id: %.0f", req.nc.connId, req.nc.conn.RemoteAddr(), req.Method, params, req.Id)
		}
		if (req.Jsonrpc != "2.0" && req.Jsonrpc != "") || req.Method == "" { //"jsonrpc":"2.0" is optional
			req.Error(ErrInvalidRequest, "", nil)
			continue
		}
		go nc.handleReq(req)
	}
}
<|endoftext|>


<commit_before>package consul

import (
	"container/list"
	"crypto/tls"
	"fmt"
	"io"
	"net"
	"net/rpc"
	"sync"
	"sync/atomic"
	"time"

	"github.com/hashicorp/consul/tlsutil"
	"github.com/hashicorp/go-msgpack/codec"
	"github.com/hashicorp/yamux"
	"github.com/inconshreveable/muxado"
)

// msgpackHandle is a shared handle for encoding/decoding of RPC messages
var msgpackHandle = &codec.MsgpackHandle{}

// muxSession is used to provide an interface for either muxado or yamux
type muxSession interface {
	Open() (net.Conn, error)
	Close() error
}

type muxadoWrapper struct {
	m muxado.Session
}

func (w *muxadoWrapper) Open() (net.Conn, error) {
	return w.m.Open()
}

func (w *muxadoWrapper) Close() error {
	return w.m.Close()
}

// streamClient is used to wrap a stream with an RPC client
type StreamClient struct {
	stream net.Conn
	client *rpc.Client
}

// Conn is a pooled connection to a Consul server
type Conn struct {
	refCount    int32
	shouldClose int32

	addr     net.Addr
	session  muxSession
	lastUsed time.Time
	version  int

	pool *ConnPool

	clients    *list.List
	clientLock sync.Mutex
}

func (c *Conn) Close() error {
	return c.session.Close()
}

// getClient is used to get a cached or new client
func (c *Conn) getClient() (*StreamClient, error) {
	// Check for cached client
	c.clientLock.Lock()
	front := c.clients.Front()
	if front != nil {
		c.clients.Remove(front)
	}
	c.clientLock.Unlock()
	if front != nil {
		return front.Value.(*StreamClient), nil
	}

	// Open a new session
	stream, err := c.session.Open()
	if err != nil {
		return nil, err
	}

	// Create the RPC client
	cc := codec.GoRpc.ClientCodec(stream, msgpackHandle)
	client := rpc.NewClientWithCodec(cc)

	// Return a new stream client
	sc := &StreamClient{
		stream: stream,
		client: client,
	}
	return sc, nil
}

// returnStream is used when done with a stream
// to allow re-use by a future RPC
func (c *Conn) returnClient(client *StreamClient) {
	didSave := false
	c.clientLock.Lock()
	if c.clients.Len() < c.pool.maxStreams && atomic.LoadInt32(&c.shouldClose) == 0 {
		c.clients.PushFront(client)
		didSave = true
	}
	c.clientLock.Unlock()
	if !didSave {
		client.stream.Close()
	}
}

// ConnPool is used to maintain a connection pool to other
// Consul servers. This is used to reduce the latency of
// RPC requests between servers. It is only used to pool
// connections in the rpcConsul mode. Raft connections
// are pooled separately.
type ConnPool struct {
	sync.Mutex

	// LogOutput is used to control logging
	logOutput io.Writer

	// The maximum time to keep a connection open
	maxTime time.Duration

	// The maximum number of open streams to keep
	maxStreams int

	// Pool maps an address to a open connection
	pool map[string]*Conn

	// TLS settings
	tlsConfig *tls.Config

	// Used to indicate the pool is shutdown
	shutdown   bool
	shutdownCh chan struct{}
}

// NewPool is used to make a new connection pool
// Maintain at most one connection per host, for up to maxTime.
// Set maxTime to 0 to disable reaping. maxStreams is used to control
// the number of idle streams allowed.
// If TLS settings are provided outgoing connections use TLS.
func NewPool(logOutput io.Writer, maxTime time.Duration, maxStreams int, tlsConfig *tls.Config) *ConnPool {
	pool := &ConnPool{
		logOutput:  logOutput,
		maxTime:    maxTime,
		maxStreams: maxStreams,
		pool:       make(map[string]*Conn),
		tlsConfig:  tlsConfig,
		shutdownCh: make(chan struct{}),
	}
	if maxTime > 0 {
		go pool.reap()
	}
	return pool
}

// Shutdown is used to close the connection pool
func (p *ConnPool) Shutdown() error {
	p.Lock()
	defer p.Unlock()

	for _, conn := range p.pool {
		conn.Close()
	}
	p.pool = make(map[string]*Conn)

	if p.shutdown {
		return nil
	}
	p.shutdown = true
	close(p.shutdownCh)
	return nil
}

// Acquire is used to get a connection that is
// pooled or to return a new connection
func (p *ConnPool) acquire(addr net.Addr, version int) (*Conn, error) {
	// Check for a pooled ocnn
	if conn := p.getPooled(addr, version); conn != nil {
		return conn, nil
	}

	// Create a new connection
	return p.getNewConn(addr, version)
}

// getPooled is used to return a pooled connection
func (p *ConnPool) getPooled(addr net.Addr, version int) *Conn {
	p.Lock()
	c := p.pool[addr.String()]
	if c != nil {
		c.lastUsed = time.Now()
		atomic.AddInt32(&c.refCount, 1)
	}
	p.Unlock()
	return c
}

// getNewConn is used to return a new connection
func (p *ConnPool) getNewConn(addr net.Addr, version int) (*Conn, error) {
	// Try to dial the conn
	conn, err := net.DialTimeout("tcp", addr.String(), 10*time.Second)
	if err != nil {
		return nil, err
	}

	// Cast to TCPConn
	if tcp, ok := conn.(*net.TCPConn); ok {
		tcp.SetKeepAlive(true)
		tcp.SetNoDelay(true)
	}

	// Check if TLS is enabled
	if p.tlsConfig != nil {
		// Switch the connection into TLS mode
		if _, err := conn.Write([]byte{byte(rpcTLS)}); err != nil {
			conn.Close()
			return nil, err
		}

		// Wrap the connection in a TLS client
		tlsConn, err := tlsutil.WrapTLSClient(conn, p.tlsConfig)
		if err != nil {
			conn.Close()
			return nil, err
		}
		conn = tlsConn
	}

	// Switch the multiplexing based on version
	var session muxSession
	if version < 2 {
		// Write the Consul multiplex byte to set the mode
		if _, err := conn.Write([]byte{byte(rpcMultiplex)}); err != nil {
			conn.Close()
			return nil, err
		}

		// Create a multiplexed session
		session = &muxadoWrapper{muxado.Client(conn)}

	} else {
		// Write the Consul multiplex byte to set the mode
		if _, err := conn.Write([]byte{byte(rpcMultiplexV2)}); err != nil {
			conn.Close()
			return nil, err
		}

		// Setup the logger
		conf := yamux.DefaultConfig()
		conf.LogOutput = p.logOutput

		// Create a multiplexed session
		session, _ = yamux.Client(conn, conf)
	}

	// Wrap the connection
	c := &Conn{
		refCount: 1,
		addr:     addr,
		session:  session,
		clients:  list.New(),
		lastUsed: time.Now(),
		version:  version,
		pool:     p,
	}

	// Track this connection, handle potential race condition
	p.Lock()
	if existing := p.pool[addr.String()]; existing != nil {
		c.Close()
		p.Unlock()
		return existing, nil
	} else {
		p.pool[addr.String()] = c
		p.Unlock()
		return c, nil
	}
}

// clearConn is used to clear any cached connection, potentially in response to an erro
func (p *ConnPool) clearConn(conn *Conn) {
	// Ensure returned streams are closed
	atomic.StoreInt32(&conn.shouldClose, 1)

	// Clear from the cache
	p.Lock()
	if c, ok := p.pool[conn.addr.String()]; ok && c == conn {
		delete(p.pool, conn.addr.String())
	}
	p.Unlock()

	// Close down immediately if idle
	if refCount := atomic.LoadInt32(&conn.refCount); refCount == 0 {
		conn.Close()
	}
}

// releaseConn is invoked when we are done with a conn to reduce the ref count
func (p *ConnPool) releaseConn(conn *Conn) {
	refCount := atomic.AddInt32(&conn.refCount, -1)
	if refCount == 0 && atomic.LoadInt32(&conn.shouldClose) == 1 {
		conn.Close()
	}
}

// getClient is used to get a usable client for an address and protocol version
func (p *ConnPool) getClient(addr net.Addr, version int) (*Conn, *StreamClient, error) {
	retries := 0
START:
	// Try to get a conn first
	conn, err := p.acquire(addr, version)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to get conn: %v", err)
	}

	// Get a client
	client, err := conn.getClient()
	if err != nil {
		p.clearConn(conn)
		p.releaseConn(conn)

		// Try to redial, possible that the TCP session closed due to timeout
		if retries == 0 {
			retries++
			goto START
		}
		return nil, nil, fmt.Errorf("failed to start stream: %v", err)
	}
	return conn, client, nil
}

// RPC is used to make an RPC call to a remote host
func (p *ConnPool) RPC(addr net.Addr, version int, method string, args interface{}, reply interface{}) error {
	// Get a usable client
	conn, sc, err := p.getClient(addr, version)
	if err != nil {
		return fmt.Errorf("rpc error: %v", err)
	}

	// Make the RPC call
	err = sc.client.Call(method, args, reply)
	if err != nil {
		p.clearConn(conn)
		p.releaseConn(conn)
		return fmt.Errorf("rpc error: %v", err)
	}

	// Done with the connection
	conn.returnClient(sc)
	p.releaseConn(conn)
	return nil
}

// Reap is used to close conns open over maxTime
func (p *ConnPool) reap() {
	for {
		// Sleep for a while
		select {
		case <-p.shutdownCh:
			return
		case <-time.After(time.Second):
		}

		// Reap all old conns
		p.Lock()
		var removed []string
		now := time.Now()
		for host, conn := range p.pool {
			// Skip recently used connections
			if now.Sub(conn.lastUsed) < p.maxTime {
				continue
			}

			// Skip connections with active streams
			if atomic.LoadInt32(&conn.refCount) > 0 {
				continue
			}

			// Close the conn
			conn.Close()

			// Remove from pool
			removed = append(removed, host)
		}
		for _, host := range removed {
			delete(p.pool, host)
		}
		p.Unlock()
	}
}
<commit_msg>* use defer to avoid tracking lock * simplify control flow<commit_after>package consul

import (
	"container/list"
	"crypto/tls"
	"fmt"
	"io"
	"net"
	"net/rpc"
	"sync"
	"sync/atomic"
	"time"

	"github.com/hashicorp/consul/tlsutil"
	"github.com/hashicorp/go-msgpack/codec"
	"github.com/hashicorp/yamux"
	"github.com/inconshreveable/muxado"
)

// msgpackHandle is a shared handle for encoding/decoding of RPC messages
var msgpackHandle = &codec.MsgpackHandle{}

// muxSession is used to provide an interface for either muxado or yamux
type muxSession interface {
	Open() (net.Conn, error)
	Close() error
}

type muxadoWrapper struct {
	m muxado.Session
}

func (w *muxadoWrapper) Open() (net.Conn, error) {
	return w.m.Open()
}

func (w *muxadoWrapper) Close() error {
	return w.m.Close()
}

// streamClient is used to wrap a stream with an RPC client
type StreamClient struct {
	stream net.Conn
	client *rpc.Client
}

// Conn is a pooled connection to a Consul server
type Conn struct {
	refCount    int32
	shouldClose int32

	addr     net.Addr
	session  muxSession
	lastUsed time.Time
	version  int

	pool *ConnPool

	clients    *list.List
	clientLock sync.Mutex
}

func (c *Conn) Close() error {
	return c.session.Close()
}

// getClient is used to get a cached or new client
func (c *Conn) getClient() (*StreamClient, error) {
	// Check for cached client
	c.clientLock.Lock()
	front := c.clients.Front()
	if front != nil {
		c.clients.Remove(front)
	}
	c.clientLock.Unlock()
	if front != nil {
		return front.Value.(*StreamClient), nil
	}

	// Open a new session
	stream, err := c.session.Open()
	if err != nil {
		return nil, err
	}

	// Create the RPC client
	cc := codec.GoRpc.ClientCodec(stream, msgpackHandle)
	client := rpc.NewClientWithCodec(cc)

	// Return a new stream client
	sc := &StreamClient{
		stream: stream,
		client: client,
	}
	return sc, nil
}

// returnStream is used when done with a stream
// to allow re-use by a future RPC
func (c *Conn) returnClient(client *StreamClient) {
	didSave := false
	c.clientLock.Lock()
	if c.clients.Len() < c.pool.maxStreams && atomic.LoadInt32(&c.shouldClose) == 0 {
		c.clients.PushFront(client)
		didSave = true
	}
	c.clientLock.Unlock()
	if !didSave {
		client.stream.Close()
	}
}

// ConnPool is used to maintain a connection pool to other
// Consul servers. This is used to reduce the latency of
// RPC requests between servers. It is only used to pool
// connections in the rpcConsul mode. Raft connections
// are pooled separately.
type ConnPool struct {
	sync.Mutex

	// LogOutput is used to control logging
	logOutput io.Writer

	// The maximum time to keep a connection open
	maxTime time.Duration

	// The maximum number of open streams to keep
	maxStreams int

	// Pool maps an address to a open connection
	pool map[string]*Conn

	// TLS settings
	tlsConfig *tls.Config

	// Used to indicate the pool is shutdown
	shutdown   bool
	shutdownCh chan struct{}
}

// NewPool is used to make a new connection pool
// Maintain at most one connection per host, for up to maxTime.
// Set maxTime to 0 to disable reaping. maxStreams is used to control
// the number of idle streams allowed.
// If TLS settings are provided outgoing connections use TLS.
func NewPool(logOutput io.Writer, maxTime time.Duration, maxStreams int, tlsConfig *tls.Config) *ConnPool {
	pool := &ConnPool{
		logOutput:  logOutput,
		maxTime:    maxTime,
		maxStreams: maxStreams,
		pool:       make(map[string]*Conn),
		tlsConfig:  tlsConfig,
		shutdownCh: make(chan struct{}),
	}
	if maxTime > 0 {
		go pool.reap()
	}
	return pool
}

// Shutdown is used to close the connection pool
func (p *ConnPool) Shutdown() error {
	p.Lock()
	defer p.Unlock()

	for _, conn := range p.pool {
		conn.Close()
	}
	p.pool = make(map[string]*Conn)

	if p.shutdown {
		return nil
	}
	p.shutdown = true
	close(p.shutdownCh)
	return nil
}

// Acquire is used to get a connection that is
// pooled or to return a new connection
func (p *ConnPool) acquire(addr net.Addr, version int) (*Conn, error) {
	// Check for a pooled ocnn
	if conn := p.getPooled(addr, version); conn != nil {
		return conn, nil
	}

	// Create a new connection
	return p.getNewConn(addr, version)
}

// getPooled is used to return a pooled connection
func (p *ConnPool) getPooled(addr net.Addr, version int) *Conn {
	p.Lock()
	c := p.pool[addr.String()]
	if c != nil {
		c.lastUsed = time.Now()
		atomic.AddInt32(&c.refCount, 1)
	}
	p.Unlock()
	return c
}

// getNewConn is used to return a new connection
func (p *ConnPool) getNewConn(addr net.Addr, version int) (*Conn, error) {
	// Try to dial the conn
	conn, err := net.DialTimeout("tcp", addr.String(), 10*time.Second)
	if err != nil {
		return nil, err
	}

	// Cast to TCPConn
	if tcp, ok := conn.(*net.TCPConn); ok {
		tcp.SetKeepAlive(true)
		tcp.SetNoDelay(true)
	}

	// Check if TLS is enabled
	if p.tlsConfig != nil {
		// Switch the connection into TLS mode
		if _, err := conn.Write([]byte{byte(rpcTLS)}); err != nil {
			conn.Close()
			return nil, err
		}

		// Wrap the connection in a TLS client
		tlsConn, err := tlsutil.WrapTLSClient(conn, p.tlsConfig)
		if err != nil {
			conn.Close()
			return nil, err
		}
		conn = tlsConn
	}

	// Switch the multiplexing based on version
	var session muxSession
	if version < 2 {
		// Write the Consul multiplex byte to set the mode
		if _, err := conn.Write([]byte{byte(rpcMultiplex)}); err != nil {
			conn.Close()
			return nil, err
		}

		// Create a multiplexed session
		session = &muxadoWrapper{muxado.Client(conn)}

	} else {
		// Write the Consul multiplex byte to set the mode
		if _, err := conn.Write([]byte{byte(rpcMultiplexV2)}); err != nil {
			conn.Close()
			return nil, err
		}

		// Setup the logger
		conf := yamux.DefaultConfig()
		conf.LogOutput = p.logOutput

		// Create a multiplexed session
		session, _ = yamux.Client(conn, conf)
	}

	// Wrap the connection
	var c *Conn

	// Track this connection, handle potential race condition
	p.Lock()
	defer p.Unlock()

	if existing := p.pool[addr.String()]; existing != nil {
		c = existing
	} else {
		c = &Conn{
			refCount: 1,
			addr:     addr,
			session:  session,
			clients:  list.New(),
			lastUsed: time.Now(),
			version:  version,
			pool:     p,
		}
		p.pool[addr.String()] = c
	}

	return c, nil
}

// clearConn is used to clear any cached connection, potentially in response to an erro
func (p *ConnPool) clearConn(conn *Conn) {
	// Ensure returned streams are closed
	atomic.StoreInt32(&conn.shouldClose, 1)

	// Clear from the cache
	p.Lock()
	if c, ok := p.pool[conn.addr.String()]; ok && c == conn {
		delete(p.pool, conn.addr.String())
	}
	p.Unlock()

	// Close down immediately if idle
	if refCount := atomic.LoadInt32(&conn.refCount); refCount == 0 {
		conn.Close()
	}
}

// releaseConn is invoked when we are done with a conn to reduce the ref count
func (p *ConnPool) releaseConn(conn *Conn) {
	refCount := atomic.AddInt32(&conn.refCount, -1)
	if refCount == 0 && atomic.LoadInt32(&conn.shouldClose) == 1 {
		conn.Close()
	}
}

// getClient is used to get a usable client for an address and protocol version
func (p *ConnPool) getClient(addr net.Addr, version int) (*Conn, *StreamClient, error) {
	retries := 0
START:
	// Try to get a conn first
	conn, err := p.acquire(addr, version)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to get conn: %v", err)
	}

	// Get a client
	client, err := conn.getClient()
	if err != nil {
		p.clearConn(conn)
		p.releaseConn(conn)

		// Try to redial, possible that the TCP session closed due to timeout
		if retries == 0 {
			retries++
			goto START
		}
		return nil, nil, fmt.Errorf("failed to start stream: %v", err)
	}
	return conn, client, nil
}

// RPC is used to make an RPC call to a remote host
func (p *ConnPool) RPC(addr net.Addr, version int, method string, args interface{}, reply interface{}) error {
	// Get a usable client
	conn, sc, err := p.getClient(addr, version)
	if err != nil {
		return fmt.Errorf("rpc error: %v", err)
	}

	// Make the RPC call
	err = sc.client.Call(method, args, reply)
	if err != nil {
		p.clearConn(conn)
		p.releaseConn(conn)
		return fmt.Errorf("rpc error: %v", err)
	}

	// Done with the connection
	conn.returnClient(sc)
	p.releaseConn(conn)
	return nil
}

// Reap is used to close conns open over maxTime
func (p *ConnPool) reap() {
	for {
		// Sleep for a while
		select {
		case <-p.shutdownCh:
			return
		case <-time.After(time.Second):
		}

		// Reap all old conns
		p.Lock()
		var removed []string
		now := time.Now()
		for host, conn := range p.pool {
			// Skip recently used connections
			if now.Sub(conn.lastUsed) < p.maxTime {
				continue
			}

			// Skip connections with active streams
			if atomic.LoadInt32(&conn.refCount) > 0 {
				continue
			}

			// Close the conn
			conn.Close()

			// Remove from pool
			removed = append(removed, host)
		}
		for _, host := range removed {
			delete(p.pool, host)
		}
		p.Unlock()
	}
}
<|endoftext|>


<commit_before>package uploads

import (
	"os"

	"github.com/materials-commons/gohandy/file"
	"github.com/materials-commons/mcstore/pkg/app"
	"github.com/materials-commons/mcstore/pkg/db/dai"
	"github.com/materials-commons/mcstore/pkg/db/schema"
	"github.com/materials-commons/mcstore/pkg/files"
	"github.com/materials-commons/mcstore/server/mcstore/uploads/processor"
)

// A finisher finishes the rest of the book keeping
// when a file has been successfully uploaded and
// reconstructed.
type finisher struct {
	files dai.Files
	dirs  dai.Dirs
	fops  file.Operations
}

// newFinisher creates a new finisher.
func newFinisher(files dai.Files, dirs dai.Dirs) *finisher {
	return &finisher{
		files: files,
		dirs:  dirs,
		fops:  file.OS,
	}
}

// finish takes care of updating the file and directory pointers, determining
// if a matching file (by checksum) has already been uploaded, and making the
// file ready for the user to access.
// TODO: refactor this method into a few separate methods that contain the
// logical blocks.
func (f *finisher) finish(req *UploadRequest, fileID, checksum string, upload *schema.Upload) error {
	filePath := app.MCDir.FilePath(fileID)

	parentID, err := f.parentID(upload.File.Name, upload.DirectoryID)
	if err != nil {
		app.Log.Errorf("Looking up parentID for %s/%s returned unexpected error: %s.", upload.File.Name, upload.DirectoryID, err)
		return err
	}

	size := f.size(fileID)

	if size != req.FlowTotalSize {
		app.Log.Errorf("Uploaded file (%s/%s) doesn't match the expected size. Expected:%d, Got: %d", upload.File.Name, req.FlowIdentifier, req.FlowTotalSize, size)
		return app.ErrInvalid
	}

	mediatype := files.MediaType(upload.File.Name, filePath)
	fields := map[string]interface{}{
		schema.FileFields.Current():   true,
		schema.FileFields.Parent():    parentID,
		schema.FileFields.Uploaded():  req.FlowTotalSize,
		schema.FileFields.Size():      req.FlowTotalSize,
		schema.FileFields.Checksum():  checksum,
		schema.FileFields.MediaType(): mediatype,
	}

	matchingFile, err := f.files.ByChecksum(checksum)
	switch {
	case err != nil && err == app.ErrNotFound:
		// This is a brand new upload for a file we haven't seen before. There are processing
		// steps that may need to be done on the file. For example we convert tif and bmp
		// image files so they can be displayed in the browser.
		f.processFile(fileID, mediatype)
	case err != nil:
		// Some type of error accessing the database
		app.Log.Errorf("Looking up file by checksum for %s/%s returned unexpected error: %s.", checksum, req.FlowFileName, err)
		return err
	default:
		// Found a matching checksum. There are two cases
		// 1. The existing file is the file we uploaded.
		// 2. The existing file is not the file we uploaded

		// Case 1: Is matching file. Delete it completely.
		if f.fileInDir(checksum, upload.File.Name, upload.DirectoryID) {
			app.Log.Infof("Found exact matching file (%s/%s) in same directory (%s), deleting", upload.File.Name, fileID, upload.DirectoryID)
			f.deleteUploadedFile(fileID, upload)
			return nil
		}

		// Case 2: Existing file is not the file we uploaded:
		// File found is not the file we uploaded, so we need to do a couple of things:
		// First set file entry to point this file (set UsesID) and secondly
		// remove the file (on the file system) that was just uploaded since we don't
		// keep duplicates.

		fields[schema.FileFields.UsesID()] = matchingFile.ID
		os.Remove(filePath)
	}

	return f.files.UpdateFields(fileID, fields)
}

// Size gets the size of the reconstructed file.
func (f *finisher) size(fileID string) int64 {
	finfo, err := os.Stat(app.MCDir.FilePath(fileID))
	if err != nil {
		return 0
	}
	return finfo.Size()
}

// parentID returns the parent for this file (if any).
func (f *finisher) parentID(fileName, dirID string) (parentID string, err error) {
	parent, err := f.files.ByPath(fileName, dirID)

	if parent != nil {
		parentID = parent.ID
	}

	if err != nil && err == app.ErrNotFound {
		// log error
		return parentID, nil
	}

	return parentID, err
}

// processFile will process the file on disk.
func (f *finisher) processFile(fileID string, mediatype schema.MediaType) {
	fp := processor.New(fileID, mediatype)
	fp.Process()
}

// fileInDir determines if this exact file has already been uploaded
// to this directory.
func (f *finisher) fileInDir(checksum, fileName, dirID string) bool {
	files, err := f.dirs.Files(dirID)
	if err != nil {
		return false
	}

	for _, fileEntry := range files {
		if fileEntry.Name == fileName && fileEntry.Checksum == checksum {
			return true
		}
	}
	return false
}

// deleteUploadedFile will completely delete the file from the system. Completely
// means it also deletes the database entries for the file.
func (f *finisher) deleteUploadedFile(fileID string, upload *schema.Upload) {
	f.files.Delete(fileID, upload.DirectoryID, upload.ProjectID)
	os.Remove(app.MCDir.FilePath(fileID))
}
<commit_msg>Use fops for Stat call so we can inject a mock for testing.<commit_after>package uploads

import (
	"os"

	"github.com/materials-commons/gohandy/file"
	"github.com/materials-commons/mcstore/pkg/app"
	"github.com/materials-commons/mcstore/pkg/db/dai"
	"github.com/materials-commons/mcstore/pkg/db/schema"
	"github.com/materials-commons/mcstore/pkg/files"
	"github.com/materials-commons/mcstore/server/mcstore/uploads/processor"
)

// A finisher finishes the rest of the book keeping
// when a file has been successfully uploaded and
// reconstructed.
type finisher struct {
	files dai.Files
	dirs  dai.Dirs
	fops  file.Operations
}

// newFinisher creates a new finisher.
func newFinisher(files dai.Files, dirs dai.Dirs) *finisher {
	return &finisher{
		files: files,
		dirs:  dirs,
		fops:  file.OS,
	}
}

// finish takes care of updating the file and directory pointers, determining
// if a matching file (by checksum) has already been uploaded, and making the
// file ready for the user to access.
func (f *finisher) finish(req *UploadRequest, fileID, checksum string, upload *schema.Upload) error {
	filePath := app.MCDir.FilePath(fileID)

	parentID, err := f.parentID(upload.File.Name, upload.DirectoryID)
	if err != nil {
		app.Log.Errorf("Looking up parentID for %s/%s returned unexpected error: %s.", upload.File.Name, upload.DirectoryID, err)
		return err
	}

	size := f.size(fileID)

	if size != req.FlowTotalSize {
		app.Log.Errorf("Uploaded file (%s/%s) doesn't match the expected size. Expected:%d, Got: %d", upload.File.Name, req.FlowIdentifier, req.FlowTotalSize, size)
		return app.ErrInvalid
	}

	mediatype := files.MediaType(upload.File.Name, filePath)
	fields := map[string]interface{}{
		schema.FileFields.Current():   true,
		schema.FileFields.Parent():    parentID,
		schema.FileFields.Uploaded():  req.FlowTotalSize,
		schema.FileFields.Size():      req.FlowTotalSize,
		schema.FileFields.Checksum():  checksum,
		schema.FileFields.MediaType(): mediatype,
	}

	matchingFile, err := f.files.ByChecksum(checksum)
	switch {
	case err != nil && err == app.ErrNotFound:
		// This is a brand new upload for a file we haven't seen before. There are processing
		// steps that may need to be done on the file. For example we convert tif and bmp
		// image files so they can be displayed in the browser.
		f.processFile(fileID, mediatype)
	case err != nil:
		// Some type of error accessing the database
		app.Log.Errorf("Looking up file by checksum for %s/%s returned unexpected error: %s.", checksum, req.FlowFileName, err)
		return err
	default:
		// Found a matching checksum. There are two cases
		// 1. The existing file is the file we uploaded.
		// 2. The existing file is not the file we uploaded

		// Case 1: Is matching file. Delete it completely.
		if f.fileInDir(checksum, upload.File.Name, upload.DirectoryID) {
			app.Log.Infof("Found exact matching file (%s/%s) in same directory (%s), deleting", upload.File.Name, fileID, upload.DirectoryID)
			f.deleteUploadedFile(fileID, upload)
			return nil
		}

		// Case 2: Existing file is not the file we uploaded:
		// File found is not the file we uploaded, so we need to do a couple of things:
		// First set file entry to point this file (set UsesID) and secondly
		// remove the file (on the file system) that was just uploaded since we don't
		// keep duplicates.

		fields[schema.FileFields.UsesID()] = matchingFile.ID
		os.Remove(filePath)
	}

	return f.files.UpdateFields(fileID, fields)
}

// Size gets the size of the reconstructed file.
func (f *finisher) size(fileID string) int64 {
	finfo, err := f.fops.Stat(app.MCDir.FilePath(fileID))
	if err != nil {
		return 0
	}
	return finfo.Size()
}

// parentID returns the parent for this file (if any).
func (f *finisher) parentID(fileName, dirID string) (parentID string, err error) {
	parent, err := f.files.ByPath(fileName, dirID)

	if parent != nil {
		parentID = parent.ID
	}

	if err != nil && err == app.ErrNotFound {
		// log error
		return parentID, nil
	}

	return parentID, err
}

// processFile will process the file on disk.
func (f *finisher) processFile(fileID string, mediatype schema.MediaType) {
	fp := processor.New(fileID, mediatype)
	fp.Process()
}

// fileInDir determines if this exact file has already been uploaded
// to this directory.
func (f *finisher) fileInDir(checksum, fileName, dirID string) bool {
	files, err := f.dirs.Files(dirID)
	if err != nil {
		return false
	}

	for _, fileEntry := range files {
		if fileEntry.Name == fileName && fileEntry.Checksum == checksum {
			return true
		}
	}
	return false
}

// deleteUploadedFile will completely delete the file from the system. Completely
// means it also deletes the database entries for the file.
func (f *finisher) deleteUploadedFile(fileID string, upload *schema.Upload) {
	f.files.Delete(fileID, upload.DirectoryID, upload.ProjectID)
	os.Remove(app.MCDir.FilePath(fileID))
}
<|endoftext|>


<commit_before>package uploads

import (
	"path/filepath"

	"github.com/materials-commons/gohandy/file"
	"github.com/materials-commons/mcstore/pkg/app"
	"github.com/materials-commons/mcstore/pkg/app/flow"
	"github.com/materials-commons/mcstore/pkg/db"
	"github.com/materials-commons/mcstore/pkg/db/dai"
	"github.com/materials-commons/mcstore/pkg/db/schema"
)

// A UploadRequest contains the block to upload and the
// information required to write that block.
type UploadRequest struct {
	*flow.Request
}

// UploadService takes care of uploading blocks and constructing the
// file when all blocks have been uploaded.
type UploadService interface {
	Upload(req *UploadRequest) error
}

// uploadService is an implementation of UploadService.
type uploadService struct {
	tracker     tracker
	files       dai.Files
	uploads     dai.Uploads
	dirs        dai.Dirs
	writer      requestWriter
	requestPath requestPath
	fops        file.Operations
}

// NewUploadService creates a new uploadService connecting
// to the default database. It will panic if it cannot
// establish a connection to the database.
func NewUploadService() *uploadService {
	session := db.RSessionMust()
	return &uploadService{
		tracker:     requestBlockTracker,
		files:       dai.NewRFiles(session),
		uploads:     dai.NewRUploads(session),
		dirs:        dai.NewRDirs(session),
		writer:      &blockRequestWriter{},
		requestPath: &mcdirRequestPath{},
		fops:        file.OS,
	}
}

// Upload performs uploading a block and constructing the file
// after all blocks have been uploaded.
func (s *uploadService) Upload(req *UploadRequest) error {
	dir := s.requestPath.dir(req.Request)
	if err := s.writer.write(dir, req.Request); err != nil {
		return err
	}

	id := req.UploadID()
	s.tracker.setBlock(id, int(req.FlowChunkNumber))

	if s.tracker.done(id) {
		// assembling the file, and performing any processing may take a while, especially
		// on large uploads, so we perform this step in the background.
		if file, err := s.assemble(req, dir); err != nil {
			app.Log.Errorf("Assembly failed for request %s: %s", req.FlowIdentifier, err)
			// Assembly failed. If file isn't nil then
			// there is some cleanup to do in the database.
			if file != nil {
				if err := s.cleanup(req, file.ID); err != nil {
					app.Log.Errorf("Attempted cleanup of failed assembly %s errored with: %s", req.FlowIdentifier, err)
				}
			}
			return err
		}
	}
	return nil
}

// assemble put the chunks for the file back together, create a database entry
// and take care of all book keeping tasks to make the file accessible.
func (s *uploadService) assemble(req *UploadRequest, dir string) (*schema.File, error) {
	// Look up the upload
	upload, err := s.uploads.ByID(req.FlowIdentifier)
	if err != nil {
		return nil, err
	}

	// Create file entry in database
	file, err := s.createFile(req, upload)
	if err != nil {
		app.Log.Errorf("Assembly failed for request %s, couldn't create file in database: %s", req.FlowIdentifier, err)
		return nil, err
	}

	// Create on disk entry to write chunks to
	if err := s.createDest(file.ID); err != nil {
		app.Log.Errorf("Assembly failed for request %s, couldn't create file on disk: %s", req.FlowIdentifier, err)
		return file, err
	}

	// Move file
	uploadDir := s.requestPath.dir(req.Request)
	s.fops.Rename(filepath.Join(uploadDir, req.UploadID()), app.MCDir.FilePath(file.ID))

	// Finish updating the file state.
	finisher := newFinisher(s.files, s.dirs)
	if err := finisher.finish(req, file.ID, upload); err != nil {
		app.Log.Errorf("Assembly failed for request %s, couldn't finish request: %s", req.FlowIdentifier, err)
		return file, err
	}

	app.Log.Infof("successfully upload fileID %s", file.ID)

	s.cleanupUploadRequest(req.UploadID())
	return nil, nil
}

// createFile creates the database file entry.
func (s *uploadService) createFile(req *UploadRequest, upload *schema.Upload) (*schema.File, error) {
	file := schema.NewFile(upload.File.Name, upload.ProjectOwner)

	f, err := s.files.Insert(&file, upload.DirectoryID, upload.ProjectID)
	app.Log.Infof("Created file %s, in %s %s", f.ID, upload.DirectoryID, upload.ProjectID)
	if err != nil {
		return nil, err
	}
	return f, nil
}

// createDest creates the destination file and ensures that the directory
// path is also created.
func (s *uploadService) createDest(fileID string) error {
	dir := app.MCDir.FileDir(fileID)
	if err := s.fops.MkdirAll(dir, 0700); err != nil {
		return err
	}
	return nil
}

// cleanup is called when an error has occurred. It attempts to clean up
// the state in the database for this particular entry.
func (s *uploadService) cleanup(req *UploadRequest, fileID string) error {
	upload, err := s.uploads.ByID(req.FlowIdentifier)
	if err != nil {
		return err
	}
	_, err = s.files.Delete(fileID, upload.DirectoryID, upload.ProjectID)
	return err
}

//cleanupUploadRequest removes the upload request and file chunks.
func (s *uploadService) cleanupUploadRequest(uploadID string) {
	s.tracker.clear(uploadID)
	s.uploads.Delete(uploadID)
	s.fops.RemoveAll(app.MCDir.UploadDir(uploadID))
}
<commit_msg>Compute and add to hash in the service.<commit_after>package uploads

import (
	"path/filepath"

	"github.com/materials-commons/gohandy/file"
	"github.com/materials-commons/mcstore/pkg/app"
	"github.com/materials-commons/mcstore/pkg/app/flow"
	"github.com/materials-commons/mcstore/pkg/db"
	"github.com/materials-commons/mcstore/pkg/db/dai"
	"github.com/materials-commons/mcstore/pkg/db/schema"
)

// A UploadRequest contains the block to upload and the
// information required to write that block.
type UploadRequest struct {
	*flow.Request
}

// UploadService takes care of uploading blocks and constructing the
// file when all blocks have been uploaded.
type UploadService interface {
	Upload(req *UploadRequest) error
}

// uploadService is an implementation of UploadService.
type uploadService struct {
	tracker     tracker
	files       dai.Files
	uploads     dai.Uploads
	dirs        dai.Dirs
	writer      requestWriter
	requestPath requestPath
	fops        file.Operations
}

// NewUploadService creates a new uploadService connecting
// to the default database. It will panic if it cannot
// establish a connection to the database.
func NewUploadService() *uploadService {
	session := db.RSessionMust()
	return &uploadService{
		tracker:     requestBlockTracker,
		files:       dai.NewRFiles(session),
		uploads:     dai.NewRUploads(session),
		dirs:        dai.NewRDirs(session),
		writer:      &blockRequestWriter{},
		requestPath: &mcdirRequestPath{},
		fops:        file.OS,
	}
}

// Upload performs uploading a block and constructing the file
// after all blocks have been uploaded.
func (s *uploadService) Upload(req *UploadRequest) error {
	dir := s.requestPath.dir(req.Request)
	if err := s.writer.write(dir, req.Request); err != nil {
		return err
	}

	id := req.UploadID()
	s.tracker.addToHash(id, req.Chunk)
	s.tracker.setBlock(id, int(req.FlowChunkNumber))

	if s.tracker.done(id) {
		// assembling the file, and performing any processing may take a while, especially
		// on large uploads, so we perform this step in the background.
		if file, err := s.assemble(req, dir); err != nil {
			app.Log.Errorf("Assembly failed for request %s: %s", req.FlowIdentifier, err)
			// Assembly failed. If file isn't nil then
			// there is some cleanup to do in the database.
			if file != nil {
				if err := s.cleanup(req, file.ID); err != nil {
					app.Log.Errorf("Attempted cleanup of failed assembly %s errored with: %s", req.FlowIdentifier, err)
				}
			}
			return err
		}
	}
	return nil
}

// assemble put the chunks for the file back together, create a database entry
// and take care of all book keeping tasks to make the file accessible.
func (s *uploadService) assemble(req *UploadRequest, dir string) (*schema.File, error) {
	// Look up the upload
	upload, err := s.uploads.ByID(req.FlowIdentifier)
	if err != nil {
		return nil, err
	}

	// Create file entry in database
	file, err := s.createFile(req, upload)
	if err != nil {
		app.Log.Errorf("Assembly failed for request %s, couldn't create file in database: %s", req.FlowIdentifier, err)
		return nil, err
	}

	// Create on disk entry to write chunks to
	if err := s.createDest(file.ID); err != nil {
		app.Log.Errorf("Assembly failed for request %s, couldn't create file on disk: %s", req.FlowIdentifier, err)
		return file, err
	}

	// Move file
	uploadDir := s.requestPath.dir(req.Request)
	s.fops.Rename(filepath.Join(uploadDir, req.UploadID()), app.MCDir.FilePath(file.ID))

	// Finish updating the file state.
	finisher := newFinisher(s.files, s.dirs)
	if err := finisher.finish(req, file.ID, s.tracker.hash(req.UploadID()), upload); err != nil {
		app.Log.Errorf("Assembly failed for request %s, couldn't finish request: %s", req.FlowIdentifier, err)
		return file, err
	}

	app.Log.Infof("successfully upload fileID %s", file.ID)

	s.cleanupUploadRequest(req.UploadID())
	return nil, nil
}

// createFile creates the database file entry.
func (s *uploadService) createFile(req *UploadRequest, upload *schema.Upload) (*schema.File, error) {
	file := schema.NewFile(upload.File.Name, upload.ProjectOwner)

	f, err := s.files.Insert(&file, upload.DirectoryID, upload.ProjectID)
	app.Log.Infof("Created file %s, in %s %s", f.ID, upload.DirectoryID, upload.ProjectID)
	if err != nil {
		return nil, err
	}
	return f, nil
}

// createDest creates the destination file and ensures that the directory
// path is also created.
func (s *uploadService) createDest(fileID string) error {
	dir := app.MCDir.FileDir(fileID)
	if err := s.fops.MkdirAll(dir, 0700); err != nil {
		return err
	}
	return nil
}

// cleanup is called when an error has occurred. It attempts to clean up
// the state in the database for this particular entry.
func (s *uploadService) cleanup(req *UploadRequest, fileID string) error {
	upload, err := s.uploads.ByID(req.FlowIdentifier)
	if err != nil {
		return err
	}
	_, err = s.files.Delete(fileID, upload.DirectoryID, upload.ProjectID)
	return err
}

//cleanupUploadRequest removes the upload request and file chunks.
func (s *uploadService) cleanupUploadRequest(uploadID string) {
	s.tracker.clear(uploadID)
	s.uploads.Delete(uploadID)
	s.fops.RemoveAll(app.MCDir.UploadDir(uploadID))
}
<|endoftext|>


<commit_before>package qbit

import (
	"fmt"
	"github.com/stretchr/testify/assert"
	"testing"
)

func TestTable(t *testing.T) {

	table := NewTable(
		"user",
		[]Column{
			NewColumn(
				"id",
				BigInt(),
				[]Constraint{
					PrimaryKey(),
				},
			),
			NewColumn(
				"email",
				VarChar(512),
				[]Constraint{
					Unique("email"),
					NotNull(),
				},
			),
			NewColumn(
				"bio",
				Text(),
				[]Constraint{
					NotNull(),
				},
			),
			NewColumn(
				"gender",
				Char(16),
				[]Constraint{
					Default("female"),
				},
			),
			NewColumn(
				"birth_date",
				Char(16),
				[]Constraint{
					NotNull(),
				},
			),
		},
		[]Constraint{},
	)

	q := fmt.Sprintf(`CREATE TABLE user(
	id BIGINT PRIMARY KEY,
	email VARCHAR(512) UNIQUE NOT NULL,
	bio TEXT NOT NULL,
	gender CHAR(16) DEFAULT %s,
	birth_date CHAR(16) NOT NULL
);`, "`female`")

	assert.Equal(t, table.Sql(), q)

	fmt.Println(table.Sql())
}
<commit_msg>fix Unique bug in table test<commit_after>package qbit

import (
	"fmt"
	"github.com/stretchr/testify/assert"
	"testing"
)

func TestTable(t *testing.T) {

	table := NewTable(
		"user",
		[]Column{
			NewColumn(
				"id",
				BigInt(),
				[]Constraint{
					PrimaryKey(),
				},
			),
			NewColumn(
				"email",
				VarChar(512),
				[]Constraint{
					Unique(),
					NotNull(),
				},
			),
			NewColumn(
				"bio",
				Text(),
				[]Constraint{
					NotNull(),
				},
			),
			NewColumn(
				"gender",
				Char(16),
				[]Constraint{
					Default("female"),
				},
			),
			NewColumn(
				"birth_date",
				Char(16),
				[]Constraint{
					NotNull(),
				},
			),
		},
		[]Constraint{},
	)

	q := fmt.Sprintf(`CREATE TABLE user(
	id BIGINT PRIMARY KEY,
	email VARCHAR(512) UNIQUE NOT NULL,
	bio TEXT NOT NULL,
	gender CHAR(16) DEFAULT %s,
	birth_date CHAR(16) NOT NULL
);`, "`female`")

	assert.Equal(t, table.Sql(), q)

	fmt.Println(table.Sql())
}
<|endoftext|>


<commit_before>package tablebook

import (
	"testing"

	. "github.com/smartystreets/goconvey/convey"
)

func TestTable(t *testing.T) {
	Convey("Given a empty table", t, func() {
		table, _ := NewTable("table", []string{"foo", "bar", "baz"})

		Convey("It can create a table", func() {
			// duplicated columns
			_, err := NewTable("table", []string{"foo", "foo"})
			So(err, ShouldEqual, ErrColumnExists)
		})

		Convey("It has a Name", func() {
			// ok
			So(table.Name(), ShouldEqual, "table")
		})

		Convey("It has dimensions", func() {
			table.AppendRow([]interface{}{1, "foo", 2})
			table.AppendRow([]interface{}{1, "foo", 2})

			So(table.Width(), ShouldEqual, 3)
			So(table.Height(), ShouldEqual, 2)
		})

		Convey("It can get columnNames", func() {
			So(table.ColumnNames(), ShouldResemble, []string{"foo", "bar", "baz"})
		})

		Convey("It can rename columns", func() {
			// rename one
			So(table.RenameColumn("unknown", "newColumnName"), ShouldEqual, ErrNotFound)

			// rename existing column
			So(table.RenameColumn("bar", "foo"), ShouldEqual, ErrColumnExists)

			table.RenameColumn("bar", "new")
			So(table.ColumnNames(), ShouldResemble, []string{"foo", "new", "baz"})
		})

		Convey("It can get rows", func() {
			table.AppendRow([]interface{}{1, 2, 3})
			table.AppendRow([]interface{}{4, 5, 6})
			table.AppendRow([]interface{}{7, 8, 9})

			//ok
			So(table.Rows(), ShouldResemble,
				[][]interface{}{
					{1, 2, 3},
					{4, 5, 6},
					{7, 8, 9},
				},
			)

			tr1, err1 := table.Row(1000)

			So(err1, ShouldEqual, ErrInvalidDimensions)
			So(tr1, ShouldBeNil)

			tr2, _ := table.Row(1)
			So(tr2, ShouldResemble, []interface{}{4, 5, 6})
		})

		Convey("It can append rows", func() {
			// too much columns
			So(table.AppendRow([]interface{}{1, "foo", "bar", "too_much"}), ShouldEqual, ErrInvalidDimensions)

			// ok, not enough columns
			table.AppendRow([]interface{}{1})
			tr1, _ := table.Row(0)
			So(tr1, ShouldResemble, []interface{}{1, "", ""})

			// ok
			tbl2 := table.AppendRow([]interface{}{1, "foo", 2})
			tr2, _ := table.Row(1)
			So(tr2, ShouldResemble, []interface{}{1, "foo", 2})
			So(tbl2, ShouldBeNil)
		})

		Convey("It can get columns", func() {
			table.AppendRow([]interface{}{1, 2, 3})
			table.AppendRow([]interface{}{4, 5, 6})
			table.AppendRow([]interface{}{7, 8, 9})

			col1, err1 := table.Column("bar")
			col2, err2 := table.Column("unknown")

			So(col1, ShouldResemble, []interface{}{2, 5, 8})
			So(err1, ShouldBeNil)

			So(col2, ShouldBeNil)
			So(err2, ShouldEqual, ErrNotFound)
		})

		Convey("It can append columns", func() {
			table.AppendRow([]interface{}{1, 2, 3})
			table.AppendRow([]interface{}{5, 6, 7})
			table.AppendRow([]interface{}{9, 10, 11})

			// too much rows
			So(table.AppendColumn("new_column", []interface{}{1, "foo", "bar", "too_much"}), ShouldEqual, ErrInvalidDimensions)

			// column already exists
			So(table.AppendColumn("foo", nil), ShouldEqual, ErrColumnExists)

			// ok
			table.AppendColumn("new_column", []interface{}{4, 8})
			So(table.ColumnNames()[3], ShouldEqual, "new_column")

			// ok
			So(
				table.Rows(),
				ShouldResemble,
				[][]interface{}{
					{1, 2, 3, 4},
					{5, 6, 7, 8},
					{9, 10, 11, ""},
				},
			)

			// ok
			So(table.AppendRow([]interface{}{1, "foo", 2}), ShouldBeNil)
		})

		Convey("It can append dynamic columns", func() {
			table.AppendRow([]interface{}{1, 2, 3})
			table.AppendRow([]interface{}{5, 6, 7})
			table.AppendRow([]interface{}{9, 10, 11})

			table.AppendEvaluatedColumn(
				"new_dynamic_column_1",
				func(table *Table, rowIndex int, cellIndex int) interface{} {
					return "from_new_dynamic_column_1"
				},
			)

			table.AppendEvaluatedColumn(
				"new_dynamic_column_2",
				func(table *Table, rowIndex int, cellIndex int) interface{} {
					v, _ := table.Cell(rowIndex, 3)
					return v
				},
			)

			table.AppendEvaluatedColumn(
				"new_dynamic_column_3",
				func(table *Table, rowIndex int, cellIndex int) interface{} {
					v, _ := table.Cell(rowIndex, 4)
					return v
				},
			)

			targetTable, _ := NewTable("targetTable", []string{"foo", "baz", "new_dynamic_column_3"})
			targetTable.Take([]*Table{table})

			// column already exists
			So(table.AppendEvaluatedColumn(
				"foo",
				func(table *Table, rowIndex int, cellIndex int) interface{} {
					return nil
				},
			), ShouldEqual, ErrColumnExists)

			// ok
			So(table.ColumnNames(), ShouldResemble, []string{"foo", "bar", "baz", "new_dynamic_column_1", "new_dynamic_column_2", "new_dynamic_column_3"})
			So(targetTable.ColumnNames(), ShouldResemble, []string{"foo", "baz", "new_dynamic_column_3"})
			So(
				targetTable.Rows(),
				ShouldResemble,
				[][]interface{}{
					{1, 3, "from_new_dynamic_column_1"},
					{5, 7, "from_new_dynamic_column_1"},
					{9, 11, "from_new_dynamic_column_1"},
				},
			)
		})

		Convey("It can take columns from other tables and merge it into own columns", func() {
			targetTable, _ := NewTable("targetTable", []string{"foo", "baz"})
			targetTable.AppendRow([]interface{}{"foo 1 targetTable", "baz 1 targetTable"})

			foreignTable1, _ := NewTable("foreignTable1", []string{"foo", "bar", "unknown"})
			foreignTable1.AppendRow([]interface{}{"foo 1 foreignTable1", "bar 1 foreignTable1", "baz 1 foreignTable1"})
			foreignTable1.AppendRow([]interface{}{"foo 2 foreignTable1", "bar 2 foreignTable1", "baz 2 foreignTable1"})
			foreignTable1.RenameColumn("unknown", "baz")

			foreignTable2, _ := NewTable("foreignTable2", []string{"bar", "foo"})
			foreignTable2.AppendRow([]interface{}{"bar 1 foreignTable2", "foo 1 foreignTable2"})
			foreignTable2.AppendRow([]interface{}{"bar 2 foreignTable2", "foo 2 foreignTable2"})

			targetTable.Take([]*Table{foreignTable1, foreignTable2})

			// ok
			So(
				targetTable.Rows(),
				ShouldResemble,
				[][]interface{}{
					{"foo 1 targetTable", "baz 1 targetTable"},
					{"foo 1 foreignTable1", "baz 1 foreignTable1"},
					{"foo 2 foreignTable1", "baz 2 foreignTable1"},
					{"foo 1 foreignTable2", ""},
					{"foo 2 foreignTable2", ""},
				},
			)

		})
	})
}
<commit_msg>add tests for table.cell errors<commit_after>package tablebook

import (
	"testing"

	. "github.com/smartystreets/goconvey/convey"
)

func TestTable(t *testing.T) {
	Convey("Given a empty table", t, func() {
		table, _ := NewTable("table", []string{"foo", "bar", "baz"})

		Convey("It can create a table", func() {
			// duplicated columns
			_, err := NewTable("table", []string{"foo", "foo"})
			So(err, ShouldEqual, ErrColumnExists)
		})

		Convey("It has a Name", func() {
			// ok
			So(table.Name(), ShouldEqual, "table")
		})

		Convey("It has dimensions", func() {
			table.AppendRow([]interface{}{1, "foo", 2})
			table.AppendRow([]interface{}{1, "foo", 2})

			So(table.Width(), ShouldEqual, 3)
			So(table.Height(), ShouldEqual, 2)
		})

		Convey("It can get columnNames", func() {
			So(table.ColumnNames(), ShouldResemble, []string{"foo", "bar", "baz"})
		})

		Convey("It can rename columns", func() {
			// rename one
			So(table.RenameColumn("unknown", "newColumnName"), ShouldEqual, ErrNotFound)

			// rename existing column
			So(table.RenameColumn("bar", "foo"), ShouldEqual, ErrColumnExists)

			table.RenameColumn("bar", "new")
			So(table.ColumnNames(), ShouldResemble, []string{"foo", "new", "baz"})
		})

		Convey("It can get rows", func() {
			table.AppendRow([]interface{}{1, 2, 3})
			table.AppendRow([]interface{}{4, 5, 6})
			table.AppendRow([]interface{}{7, 8, 9})

			//ok
			So(table.Rows(), ShouldResemble,
				[][]interface{}{
					{1, 2, 3},
					{4, 5, 6},
					{7, 8, 9},
				},
			)

			tr1, err1 := table.Row(1000)

			So(err1, ShouldEqual, ErrInvalidDimensions)
			So(tr1, ShouldBeNil)

			tr2, _ := table.Row(1)
			So(tr2, ShouldResemble, []interface{}{4, 5, 6})
		})

		Convey("It can append rows", func() {
			// too much columns
			So(table.AppendRow([]interface{}{1, "foo", "bar", "too_much"}), ShouldEqual, ErrInvalidDimensions)

			// ok, not enough columns
			table.AppendRow([]interface{}{1})
			tr1, _ := table.Row(0)
			So(tr1, ShouldResemble, []interface{}{1, "", ""})

			// ok
			tbl2 := table.AppendRow([]interface{}{1, "foo", 2})
			tr2, _ := table.Row(1)
			So(tr2, ShouldResemble, []interface{}{1, "foo", 2})
			So(tbl2, ShouldBeNil)
		})

		Convey("It can get columns", func() {
			table.AppendRow([]interface{}{1, 2, 3})
			table.AppendRow([]interface{}{4, 5, 6})
			table.AppendRow([]interface{}{7, 8, 9})

			col1, err1 := table.Column("bar")
			col2, err2 := table.Column("unknown")

			So(col1, ShouldResemble, []interface{}{2, 5, 8})
			So(err1, ShouldBeNil)

			So(col2, ShouldBeNil)
			So(err2, ShouldEqual, ErrNotFound)
		})

		Convey("It can append columns", func() {
			table.AppendRow([]interface{}{1, 2, 3})
			table.AppendRow([]interface{}{5, 6, 7})
			table.AppendRow([]interface{}{9, 10, 11})

			// too much rows
			So(table.AppendColumn("new_column", []interface{}{1, "foo", "bar", "too_much"}), ShouldEqual, ErrInvalidDimensions)

			// column already exists
			So(table.AppendColumn("foo", nil), ShouldEqual, ErrColumnExists)

			// ok
			table.AppendColumn("new_column", []interface{}{4, 8})
			So(table.ColumnNames()[3], ShouldEqual, "new_column")

			// ok
			So(
				table.Rows(),
				ShouldResemble,
				[][]interface{}{
					{1, 2, 3, 4},
					{5, 6, 7, 8},
					{9, 10, 11, ""},
				},
			)

			// ok
			So(table.AppendRow([]interface{}{1, "foo", 2}), ShouldBeNil)
		})

		Convey("It can return cells", func() {
			table.AppendRow([]interface{}{1, 2, 3})
			table.AppendRow([]interface{}{5, 6, 7})
			table.AppendRow([]interface{}{9, 10, 11})

			// cell out of range
			_, err1 := table.Cell(100, 100)
			_, err2 := table.Cell(-100, -100)
			So(err1, ShouldEqual, ErrInvalidDimensions)
			So(err2, ShouldEqual, ErrInvalidDimensions)

		})

		Convey("It can append dynamic columns", func() {
			table.AppendRow([]interface{}{1, 2, 3})
			table.AppendRow([]interface{}{5, 6, 7})
			table.AppendRow([]interface{}{9, 10, 11})

			table.AppendEvaluatedColumn(
				"new_dynamic_column_1",
				func(table *Table, rowIndex int, cellIndex int) interface{} {
					return "from_new_dynamic_column_1"
				},
			)

			table.AppendEvaluatedColumn(
				"new_dynamic_column_2",
				func(table *Table, rowIndex int, cellIndex int) interface{} {
					v, _ := table.Cell(rowIndex, 3)
					return v
				},
			)

			table.AppendEvaluatedColumn(
				"new_dynamic_column_3",
				func(table *Table, rowIndex int, cellIndex int) interface{} {
					v, _ := table.Cell(rowIndex, 4)
					return v
				},
			)

			targetTable, _ := NewTable("targetTable", []string{"foo", "baz", "new_dynamic_column_3"})
			targetTable.Take([]*Table{table})

			// column already exists
			So(table.AppendEvaluatedColumn(
				"foo",
				func(table *Table, rowIndex int, cellIndex int) interface{} {
					return nil
				},
			), ShouldEqual, ErrColumnExists)

			// ok
			So(table.ColumnNames(), ShouldResemble, []string{"foo", "bar", "baz", "new_dynamic_column_1", "new_dynamic_column_2", "new_dynamic_column_3"})
			So(targetTable.ColumnNames(), ShouldResemble, []string{"foo", "baz", "new_dynamic_column_3"})
			So(
				targetTable.Rows(),
				ShouldResemble,
				[][]interface{}{
					{1, 3, "from_new_dynamic_column_1"},
					{5, 7, "from_new_dynamic_column_1"},
					{9, 11, "from_new_dynamic_column_1"},
				},
			)
		})

		Convey("It can take columns from other tables and merge it into own columns", func() {
			targetTable, _ := NewTable("targetTable", []string{"foo", "baz"})
			targetTable.AppendRow([]interface{}{"foo 1 targetTable", "baz 1 targetTable"})

			foreignTable1, _ := NewTable("foreignTable1", []string{"foo", "bar", "unknown"})
			foreignTable1.AppendRow([]interface{}{"foo 1 foreignTable1", "bar 1 foreignTable1", "baz 1 foreignTable1"})
			foreignTable1.AppendRow([]interface{}{"foo 2 foreignTable1", "bar 2 foreignTable1", "baz 2 foreignTable1"})
			foreignTable1.RenameColumn("unknown", "baz")

			foreignTable2, _ := NewTable("foreignTable2", []string{"bar", "foo"})
			foreignTable2.AppendRow([]interface{}{"bar 1 foreignTable2", "foo 1 foreignTable2"})
			foreignTable2.AppendRow([]interface{}{"bar 2 foreignTable2", "foo 2 foreignTable2"})

			targetTable.Take([]*Table{foreignTable1, foreignTable2})

			// ok
			So(
				targetTable.Rows(),
				ShouldResemble,
				[][]interface{}{
					{"foo 1 targetTable", "baz 1 targetTable"},
					{"foo 1 foreignTable1", "baz 1 foreignTable1"},
					{"foo 2 foreignTable1", "baz 2 foreignTable1"},
					{"foo 1 foreignTable2", ""},
					{"foo 2 foreignTable2", ""},
				},
			)

		})
	})
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"os"
	"strconv"

	"github.com/taskcluster/httpbackoff"
	"github.com/taskcluster/taskcluster-client-go"
	"github.com/taskcluster/taskcluster-client-go/queue"
)

type TaskStatusUpdate struct {
	Task       *TaskRun
	Status     TaskStatus
	IfStatusIn map[TaskStatus]bool
	Reason     string
}

// Enumerate task status to aid life-cycle decision making
// Use strings for benefit of simple logging/reporting
const (
	Aborted   TaskStatus = "Aborted"
	Cancelled TaskStatus = "Cancelled"
	Succeeded TaskStatus = "Succeeded"
	Failed    TaskStatus = "Failed"
	Errored   TaskStatus = "Errored"
	Claimed   TaskStatus = "Claimed"
	Reclaimed TaskStatus = "Reclaimed"
)

// TaskStatusHandler is the single point of contact
// between go routines for handling updates to the
// status of a task. When called, it will spawn a
// go routine to manage task status updates (both
// internally tracking TaskRun.Status and also
// informing the Queue about the status). Errors
// from talking to the Queue are returned on err
// channel.
func TaskStatusHandler() (request chan<- TaskStatusUpdate, err <-chan error, done chan<- bool) {
	r := make(chan TaskStatusUpdate)
	e := make(chan error)
	d := make(chan bool)

	// we'll make all these functions internal to TaskStatusHandler so that
	// they can only be called inside here, so that reading/writing to the
	// appropriate channels is the only way to trigger them, to ensure
	// proper concurrency handling

	reportException := func(task *TaskRun, reason string) error {
		ter := queue.TaskExceptionRequest{Reason: reason}
		tsr, err := Queue.ReportException(task.TaskID, strconv.FormatInt(int64(task.RunID), 10), &ter)
		if err != nil {
			log.Printf("Not able to report exception for task %v:", task.TaskID)
			log.Printf("%v", err)
			return err
		}
		task.TaskClaimResponse.Status = tsr.Status
		log.Println(task.String())
		return nil
	}

	reportFailed := func(task *TaskRun) error {
		tsr, err := Queue.ReportFailed(task.TaskID, strconv.FormatInt(int64(task.RunID), 10))
		if err != nil {
			log.Printf("Not able to report failed completion for task %v:", task.TaskID)
			log.Printf("%v", err)
			return err
		}
		task.TaskClaimResponse.Status = tsr.Status
		log.Println(task.String())
		return nil
	}

	reportCompleted := func(task *TaskRun) error {
		log.Println("Command finished successfully!")
		tsr, err := Queue.ReportCompleted(task.TaskID, strconv.FormatInt(int64(task.RunID), 10))
		if err != nil {
			log.Printf("Not able to report successful completion for task %v:", task.TaskID)
			log.Printf("%v", err)
			return err
		}
		task.TaskClaimResponse.Status = tsr.Status
		// log.Println(task.String())
		return nil
	}

	claim := func(task *TaskRun) error {
		log.Printf("Claiming task %v...", task.TaskID)
		task.TaskClaimRequest = queue.TaskClaimRequest{
			WorkerGroup: config.WorkerGroup,
			WorkerID:    config.WorkerID,
		}
		// Using the taskId and runId from the <MessageText> tag, the worker
		// must call queue.claimTask().
		tcrsp, err := Queue.ClaimTask(task.TaskID, fmt.Sprintf("%d", task.RunID), &task.TaskClaimRequest)
		// check if an error occurred...
		if err != nil {
			// If the queue.claimTask() operation fails with a 4xx error, the
			// worker must delete the messages from the Azure queue (except 401).
			switch err := err.(type) {
			case httpbackoff.BadHttpResponseCode:
				switch {
				case err.HttpResponseCode == 401:
					log.Printf("Whoops - not authorized to claim task %v, *not* deleting it from Azure queue!", task.TaskID)
				case err.HttpResponseCode/100 == 4:
					// attempt to delete, but if it fails, log and continue
					// nothing we can do, and better to return the first 4xx error
					errDelete := task.deleteFromAzure()
					if errDelete != nil {
						log.Printf("Not able to delete task %v from Azure after receiving http status code %v when claiming it.", task.TaskID, err.HttpResponseCode)
						log.Printf("%v", errDelete)
					}
				}
			}
			log.Println(task.String())
			log.Printf("%v", err)
			return err
		}
		task.TaskClaimResponse = *tcrsp
		// note we don't need to worry about a mutex here since either old
		// value or new value can be used for some crossover time, and the
		// update should be atomic
		task.Queue = queue.New(&tcclient.Credentials{
			ClientID:    tcrsp.Credentials.ClientID,
			AccessToken: tcrsp.Credentials.AccessToken,
			Certificate: tcrsp.Credentials.Certificate,
		})

		// don't report failure if this fails, as it is already logged and failure =>
		task.deleteFromAzure()
		return nil
	}

	reclaim := func(task *TaskRun) error {
		log.Printf("Reclaiming task %v...", task.TaskID)
		tcrsp, err := Queue.ReclaimTask(task.TaskID, fmt.Sprintf("%d", task.RunID))

		// check if an error occurred...
		if err != nil {
			log.Printf("%v", err)
			return err
		}

		task.TaskReclaimResponse = *tcrsp
		// note we don't need to worry about a mutex here since either old
		// value or new value can be used for some crossover time, and the
		// update should be atomic
		task.Queue = queue.New(&tcclient.Credentials{
			ClientID:    tcrsp.Credentials.ClientID,
			AccessToken: tcrsp.Credentials.AccessToken,
			Certificate: tcrsp.Credentials.Certificate,
		})
		log.Printf("Reclaimed task %v successfully.", task.TaskID)
		return nil
	}

	abort := func(task *TaskRun, reason string) error {
		log.Printf("Aborting task %v due to: %v...", task.TaskID, reason)
		task.Status = Aborted
		// TODO: need to kill running jobs! Need a go routine to track running
		// jobs, and kill them on aborts
		return reportException(task, "task-cancelled")
	}

	cancel := func(task *TaskRun, reason string) error {
		//TODO: implement cancelling of tasks
		return nil
	}

	go func() {
		for {
			select {
			case update := <-r:
				// only update if either IfStatusIn is nil
				// or it is non-nil but it has "true" value
				// for key of current status
				if update.IfStatusIn == nil || update.IfStatusIn[task.Status] {
					task := update.Task
					task.Status = update.Status
					switch update.Status {
					// Aborting is when you stop running a job you already claimed
					case Aborted:
						e <- abort(task, update.Reason)
					// Cancelling is when you decide not to run a job which you haven't yet claimed
					case Cancelled:
						e <- cancel(task, update.Reason)
					case Succeeded:
						e <- reportCompleted(task)
					case Failed:
						e <- reportFailed(task)
					case Errored:
						e <- reportException(task, update.Reason)
					case Claimed:
						e <- claim(task)
					case Reclaimed:
						e <- reclaim(task)
					default:
						log.Printf("Internal error: unknown task status: %v", update.Status)
						os.Exit(64)
					}
				} else {
					// current status is such that we shouldn't update to new
					// status, so just report that no error occurred...
					log.Printf("Not able to update status to %v - current status %v, allowed current status for update: %v", update.Status, update.Task.Status, update.IfStatusIn)
					e <- nil
				}
			case <-d:
				close(d)
				break
			}
		}
	}()
	return r, e, d
}
<commit_msg>Bug 1307383 - fix for the fix ... whoops<commit_after>package main

import (
	"fmt"
	"log"
	"os"
	"strconv"

	"github.com/taskcluster/httpbackoff"
	"github.com/taskcluster/taskcluster-client-go"
	"github.com/taskcluster/taskcluster-client-go/queue"
)

type TaskStatusUpdate struct {
	Task       *TaskRun
	Status     TaskStatus
	IfStatusIn map[TaskStatus]bool
	Reason     string
}

// Enumerate task status to aid life-cycle decision making
// Use strings for benefit of simple logging/reporting
const (
	Aborted   TaskStatus = "Aborted"
	Cancelled TaskStatus = "Cancelled"
	Succeeded TaskStatus = "Succeeded"
	Failed    TaskStatus = "Failed"
	Errored   TaskStatus = "Errored"
	Claimed   TaskStatus = "Claimed"
	Reclaimed TaskStatus = "Reclaimed"
)

// TaskStatusHandler is the single point of contact
// between go routines for handling updates to the
// status of a task. When called, it will spawn a
// go routine to manage task status updates (both
// internally tracking TaskRun.Status and also
// informing the Queue about the status). Errors
// from talking to the Queue are returned on err
// channel.
func TaskStatusHandler() (request chan<- TaskStatusUpdate, err <-chan error, done chan<- bool) {
	r := make(chan TaskStatusUpdate)
	e := make(chan error)
	d := make(chan bool)

	// we'll make all these functions internal to TaskStatusHandler so that
	// they can only be called inside here, so that reading/writing to the
	// appropriate channels is the only way to trigger them, to ensure
	// proper concurrency handling

	reportException := func(task *TaskRun, reason string) error {
		ter := queue.TaskExceptionRequest{Reason: reason}
		tsr, err := Queue.ReportException(task.TaskID, strconv.FormatInt(int64(task.RunID), 10), &ter)
		if err != nil {
			log.Printf("Not able to report exception for task %v:", task.TaskID)
			log.Printf("%v", err)
			return err
		}
		task.TaskClaimResponse.Status = tsr.Status
		log.Println(task.String())
		return nil
	}

	reportFailed := func(task *TaskRun) error {
		tsr, err := Queue.ReportFailed(task.TaskID, strconv.FormatInt(int64(task.RunID), 10))
		if err != nil {
			log.Printf("Not able to report failed completion for task %v:", task.TaskID)
			log.Printf("%v", err)
			return err
		}
		task.TaskClaimResponse.Status = tsr.Status
		log.Println(task.String())
		return nil
	}

	reportCompleted := func(task *TaskRun) error {
		log.Println("Command finished successfully!")
		tsr, err := Queue.ReportCompleted(task.TaskID, strconv.FormatInt(int64(task.RunID), 10))
		if err != nil {
			log.Printf("Not able to report successful completion for task %v:", task.TaskID)
			log.Printf("%v", err)
			return err
		}
		task.TaskClaimResponse.Status = tsr.Status
		// log.Println(task.String())
		return nil
	}

	claim := func(task *TaskRun) error {
		log.Printf("Claiming task %v...", task.TaskID)
		task.TaskClaimRequest = queue.TaskClaimRequest{
			WorkerGroup: config.WorkerGroup,
			WorkerID:    config.WorkerID,
		}
		// Using the taskId and runId from the <MessageText> tag, the worker
		// must call queue.claimTask().
		tcrsp, err := Queue.ClaimTask(task.TaskID, fmt.Sprintf("%d", task.RunID), &task.TaskClaimRequest)
		// check if an error occurred...
		if err != nil {
			// If the queue.claimTask() operation fails with a 4xx error, the
			// worker must delete the messages from the Azure queue (except 401).
			switch err := err.(type) {
			case httpbackoff.BadHttpResponseCode:
				switch {
				case err.HttpResponseCode == 401:
					log.Printf("Whoops - not authorized to claim task %v, *not* deleting it from Azure queue!", task.TaskID)
				case err.HttpResponseCode/100 == 4:
					// attempt to delete, but if it fails, log and continue
					// nothing we can do, and better to return the first 4xx error
					errDelete := task.deleteFromAzure()
					if errDelete != nil {
						log.Printf("Not able to delete task %v from Azure after receiving http status code %v when claiming it.", task.TaskID, err.HttpResponseCode)
						log.Printf("%v", errDelete)
					}
				}
			}
			log.Println(task.String())
			log.Printf("%v", err)
			return err
		}
		task.TaskClaimResponse = *tcrsp
		// note we don't need to worry about a mutex here since either old
		// value or new value can be used for some crossover time, and the
		// update should be atomic
		task.Queue = queue.New(&tcclient.Credentials{
			ClientID:    tcrsp.Credentials.ClientID,
			AccessToken: tcrsp.Credentials.AccessToken,
			Certificate: tcrsp.Credentials.Certificate,
		})

		// don't report failure if this fails, as it is already logged and failure =>
		task.deleteFromAzure()
		return nil
	}

	reclaim := func(task *TaskRun) error {
		log.Printf("Reclaiming task %v...", task.TaskID)
		tcrsp, err := Queue.ReclaimTask(task.TaskID, fmt.Sprintf("%d", task.RunID))

		// check if an error occurred...
		if err != nil {
			log.Printf("%v", err)
			return err
		}

		task.TaskReclaimResponse = *tcrsp
		// note we don't need to worry about a mutex here since either old
		// value or new value can be used for some crossover time, and the
		// update should be atomic
		task.Queue = queue.New(&tcclient.Credentials{
			ClientID:    tcrsp.Credentials.ClientID,
			AccessToken: tcrsp.Credentials.AccessToken,
			Certificate: tcrsp.Credentials.Certificate,
		})
		log.Printf("Reclaimed task %v successfully.", task.TaskID)
		return nil
	}

	abort := func(task *TaskRun, reason string) error {
		log.Printf("Aborting task %v due to: %v...", task.TaskID, reason)
		task.Status = Aborted
		// TODO: need to kill running jobs! Need a go routine to track running
		// jobs, and kill them on aborts
		return reportException(task, "task-cancelled")
	}

	cancel := func(task *TaskRun, reason string) error {
		//TODO: implement cancelling of tasks
		return nil
	}

	go func() {
		for {
			select {
			case update := <-r:
				// only update if either IfStatusIn is nil
				// or it is non-nil but it has "true" value
				// for key of current status
				// note - task could have been aborted externally, so update.Task.Status
				// may not be correct in case of task canellation, but this race
				// condition will always exist, even if we add a check for the current
				// value, so then better to fail with a 409 RequestConflict
				if update.IfStatusIn == nil || update.IfStatusIn[update.Task.Status] {
					task := update.Task
					task.Status = update.Status
					switch update.Status {
					// Aborting is when you stop running a job you already claimed
					case Aborted:
						e <- abort(task, update.Reason)
					// Cancelling is when you decide not to run a job which you haven't yet claimed
					case Cancelled:
						e <- cancel(task, update.Reason)
					case Succeeded:
						e <- reportCompleted(task)
					case Failed:
						e <- reportFailed(task)
					case Errored:
						e <- reportException(task, update.Reason)
					case Claimed:
						e <- claim(task)
					case Reclaimed:
						e <- reclaim(task)
					default:
						log.Printf("Internal error: unknown task status: %v", update.Status)
						os.Exit(64)
					}
				} else {
					// current status is such that we shouldn't update to new
					// status, so just report that no error occurred...
					log.Printf("Not able to update status to %v - current status %v, allowed current status for update: %v", update.Status, update.Task.Status, update.IfStatusIn)
					e <- nil
				}
			case <-d:
				close(d)
				break
			}
		}
	}()
	return r, e, d
}
<|endoftext|>