<commit_before>package adodb

import (
	"errors"
	"database/sql"
	"database/sql/driver"
	"fmt"
	"github.com/mattn/go-ole"
	"github.com/mattn/go-ole/oleutil"
	"math/big"
	"time"
	"unsafe"
)

func init() {
	ole.CoInitialize(0)
	sql.Register("adodb", &AdodbDriver{})
}

type AdodbDriver struct {

}

type AdodbConn struct {
	db *ole.IDispatch
}

type AdodbTx struct {
	c *AdodbConn
}

func (tx *AdodbTx) Commit() error {
	_, err := oleutil.CallMethod(tx.c.db, "CommitTrans")
	if err != nil {
		return err
	}
	return nil
}

func (tx *AdodbTx) Rollback() error {
	_, err := oleutil.CallMethod(tx.c.db, "Rollback")
	if err != nil {
		return err
	}
	return nil
}

func (c *AdodbConn) exec(cmd string) error {
	_, err := oleutil.CallMethod(c.db, "Execute", cmd)
	return err
}

func (c *AdodbConn) Begin() (driver.Tx, error) {
	_, err := oleutil.CallMethod(c.db, "BeginTrans")
	if err != nil {
		return nil, err
	}
	return &AdodbTx{c}, nil
}

func (d *AdodbDriver) Open(dsn string) (driver.Conn, error) {
	unknown, err := oleutil.CreateObject("ADODB.Connection")
	if err != nil {
		return nil, err
	}
	db, err := unknown.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.CallMethod(db, "Open", dsn)
	if err != nil {
		return nil, err
	}
	return &AdodbConn{db}, nil
}

func (c *AdodbConn) Close() error {
	_, err := oleutil.CallMethod(c.db, "Close")
	if err != nil {
		return err
	}
	c.db = nil
	return nil
}

type AdodbStmt struct {
	c  *AdodbConn
	s  *ole.IDispatch
	ps *ole.IDispatch
	b  []string
}

func (c *AdodbConn) Prepare(query string) (driver.Stmt, error) {
	unknown, err := oleutil.CreateObject("ADODB.Command")
	if err != nil {
		return nil, err
	}
	s, err := unknown.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "ActiveConnection", c.db)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "CommandText", query)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "CommandType", 1)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "Prepared", true)
	if err != nil {
		return nil, err
	}
	val, err := oleutil.GetProperty(s, "Parameters")
	if err != nil {
		return nil, err
	}
	return &AdodbStmt{c, s, val.ToIDispatch(), nil}, nil
}

func (s *AdodbStmt) Bind(bind []string) error {
	s.b = bind
	return nil
}

func (s *AdodbStmt) Close() error {
	s.s.Release()
	return nil
}

func (s *AdodbStmt) NumInput() int {
	if s.b != nil {
		return len(s.b)
	}
	_, err := oleutil.CallMethod(s.ps, "Refresh")
	if err != nil {
		return -1
	}
	val, err := oleutil.GetProperty(s.ps, "Count")
	if err != nil {
		return -1
	}
	c := int(val.Val)
	return c
}

func (s *AdodbStmt) bind(args []driver.Value) error {
	if s.b != nil {
		for i, v := range args {
			var b string = "?"
			if len(s.b) < i {
				b = s.b[i]
			}
			unknown, err := oleutil.CallMethod(s.s, "CreateParameter", b, 12, 1)
			if err != nil {
				return err
			}
			param := unknown.ToIDispatch()
			defer param.Release()
			_, err = oleutil.PutProperty(param, "Value", v)
			if err != nil {
				return err
			}
			_, err = oleutil.CallMethod(s.ps, "Append", param)
			if err != nil {
				return err
			}
		}
	} else {
		for i, v := range args {
			var varval ole.VARIANT
			varval.VT = ole.VT_I4
			varval.Val = int64(i)
			val, err := oleutil.CallMethod(s.ps, "Item", &varval)
			if err != nil {
				return err
			}
			item := val.ToIDispatch()
			defer item.Release()
			_, err = oleutil.PutProperty(item, "Value", v)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func (s *AdodbStmt) Query(args []driver.Value) (driver.Rows, error) {
	if err := s.bind(args); err != nil {
		return nil, err
	}
	rc, err := oleutil.CallMethod(s.s, "Execute")
	if err != nil {
		return nil, err
	}
	return &AdodbRows{s, rc.ToIDispatch(), -1, nil}, nil
}

func (s *AdodbStmt) Exec(args []driver.Value) (driver.Result, error) {
	if err := s.bind(args); err != nil {
		return nil, err
	}
	_, err := oleutil.CallMethod(s.s, "Execute")
	if err != nil {
		return nil, err
	}
	return driver.ResultNoRows, nil
}

type AdodbRows struct {
	s    *AdodbStmt
	rc   *ole.IDispatch
	nc   int
	cols []string
}

func (rc *AdodbRows) Close() error {
	_, err := oleutil.CallMethod(rc.rc, "Close")
	if err != nil {
		return err
	}
	return nil
}

func (rc *AdodbRows) Columns() []string {
	if rc.nc != len(rc.cols) {
		unknown, err := oleutil.GetProperty(rc.rc, "Fields")
		if err != nil {
			return []string{}
		}
		fields := unknown.ToIDispatch()
		defer fields.Release()
		val, err := oleutil.GetProperty(fields, "Count")
		if err != nil {
			return []string{}
		}
		rc.nc = int(val.Val)
		rc.cols = make([]string, rc.nc)
		for i := 0; i < rc.nc; i++ {
			var varval ole.VARIANT
			varval.VT = ole.VT_I4
			varval.Val = int64(i)
			val, err := oleutil.CallMethod(fields, "Item", &varval)
			if err != nil {
				return []string{}
			}
			item := val.ToIDispatch()
			if err != nil {
				return []string{}
			}
			name, err := oleutil.GetProperty(item, "Name")
			if err != nil {
				return []string{}
			}
			rc.cols[i] = name.ToString()
			item.Release()
		}
	}
	return rc.cols
}

func (rc *AdodbRows) Next(dest []driver.Value) error {
	_, err := oleutil.CallMethod(rc.rc, "MoveNext")
	if err != nil {
		fmt.Println(err)
		return err
	}
	unknown, err := oleutil.GetProperty(rc.rc, "EOF")
	if err != nil {
		return err
	}
	if unknown.Val != 0 {
		return errors.New("EOF")
	}
	unknown, err = oleutil.GetProperty(rc.rc, "Fields")
	if err != nil {
		return err
	}
	fields := unknown.ToIDispatch()
	defer fields.Release()
	for i := range dest {
		var varval ole.VARIANT
		varval.VT = ole.VT_I4
		varval.Val = int64(i)
		val, err := oleutil.CallMethod(fields, "Item", &varval)
		if err != nil {
			return err
		}
		field := val.ToIDispatch()
		defer field.Release()
		typ, err := oleutil.GetProperty(field, "Type")
		if err != nil {
			return err
		}
		val, err = oleutil.GetProperty(field, "Value")
		if err != nil {
			return err
		}
		field.Release()
		switch typ.Val {
		case 0: // ADEMPTY
			// TODO
		case 2: // ADSMALLINT
			dest[i] = int16(val.Val)
		case 3: // ADINTEGER
			dest[i] = int32(val.Val)
		case 4: // ADSINGLE
			dest[i] = float32(val.Val)
		case 5: // ADDOUBLE
			dest[i] = val.Val
		case 6: // ADCURRENCY
			// TODO
		case 7: // ADDATE
			// TODO
		case 8: // ADBSTR
			dest[i] = val.ToString()
		case 9: // ADIDISPATCH
			dest[i] = val.ToIDispatch()
		case 10: // ADERROR
			// TODO
		case 11: // ADBOOLEAN
			if val.Val != 0 {
				dest[i] = true
			} else {
				dest[i] = false
			}
		case 12: // ADVARIANT
			dest[i] = val
		case 13: // ADIUNKNOWN
			dest[i] = val.ToIUnknown()
		case 14: // ADDECIMAL
			// TODO
		case 16: // ADTINYINT
			dest[i] = int8(val.Val)
		case 17: // ADUNSIGNEDTINYINT
			dest[i] = uint8(val.Val)
		case 18: // ADUNSIGNEDSMALLINT
			dest[i] = uint16(val.Val)
		case 19: // ADUNSIGNEDINT
			dest[i] = uint32(val.Val)
		case 20: // ADBIGINT
			dest[i] = big.NewInt(val.Val)
		case 21: // ADUNSIGNEDBIGINT
			// TODO
		case 72: // ADGUID
			// TODO
		case 128: // ADBINARY
			sa := *(**ole.SAFEARRAY)(unsafe.Pointer(&val.Val))
			dest[i] = (*[1 << 30]byte)(unsafe.Pointer(uintptr(sa.PvData)))[0:sa.CbElements]
		case 129: // ADCHAR
			dest[i] = uint8(val.Val)
		case 130: // ADWCHAR
			dest[i] = uint16(val.Val)
		case 131: // ADNUMERIC
			dest[i] = val.Val
		case 132: // ADUSERDEFINED
			dest[i] = uintptr(val.Val)
		case 133: // ADDBDATE
			dest[i] = time.Unix(0, val.Val).UTC()
		case 134: // ADDBTIME
			dest[i] = time.Unix(0, val.Val).UTC()
		case 135: // ADDBTIMESTAMP
			dest[i] = time.Unix(0, val.Val).UTC()
		case 136: // ADCHAPTER
			dest[i] = val.ToString()
		case 200: // ADVARCHAR
			dest[i] = val.ToString()
		case 201: // ADLONGVARCHAR
			dest[i] = val.ToString()
		case 202: // ADVARWCHAR
			dest[i] = val.ToString()
		case 203: // ADLONGVARWCHAR
			dest[i] = val.ToString()
		case 204: // ADVARBINARY
			// TODO
		case 205: // ADLONGVARBINARY
			// TODO
		}
	}
	return nil
}
<commit_msg>check EOF first. close #1<commit_after>package adodb

import (
	"errors"
	"database/sql"
	"database/sql/driver"
	"fmt"
	"github.com/mattn/go-ole"
	"github.com/mattn/go-ole/oleutil"
	"math/big"
	"time"
	"unsafe"
)

func init() {
	ole.CoInitialize(0)
	sql.Register("adodb", &AdodbDriver{})
}

type AdodbDriver struct {

}

type AdodbConn struct {
	db *ole.IDispatch
}

type AdodbTx struct {
	c *AdodbConn
}

func (tx *AdodbTx) Commit() error {
	_, err := oleutil.CallMethod(tx.c.db, "CommitTrans")
	if err != nil {
		return err
	}
	return nil
}

func (tx *AdodbTx) Rollback() error {
	_, err := oleutil.CallMethod(tx.c.db, "Rollback")
	if err != nil {
		return err
	}
	return nil
}

func (c *AdodbConn) exec(cmd string) error {
	_, err := oleutil.CallMethod(c.db, "Execute", cmd)
	return err
}

func (c *AdodbConn) Begin() (driver.Tx, error) {
	_, err := oleutil.CallMethod(c.db, "BeginTrans")
	if err != nil {
		return nil, err
	}
	return &AdodbTx{c}, nil
}

func (d *AdodbDriver) Open(dsn string) (driver.Conn, error) {
	unknown, err := oleutil.CreateObject("ADODB.Connection")
	if err != nil {
		return nil, err
	}
	db, err := unknown.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.CallMethod(db, "Open", dsn)
	if err != nil {
		return nil, err
	}
	return &AdodbConn{db}, nil
}

func (c *AdodbConn) Close() error {
	_, err := oleutil.CallMethod(c.db, "Close")
	if err != nil {
		return err
	}
	c.db = nil
	return nil
}

type AdodbStmt struct {
	c  *AdodbConn
	s  *ole.IDispatch
	ps *ole.IDispatch
	b  []string
}

func (c *AdodbConn) Prepare(query string) (driver.Stmt, error) {
	unknown, err := oleutil.CreateObject("ADODB.Command")
	if err != nil {
		return nil, err
	}
	s, err := unknown.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "ActiveConnection", c.db)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "CommandText", query)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "CommandType", 1)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "Prepared", true)
	if err != nil {
		return nil, err
	}
	val, err := oleutil.GetProperty(s, "Parameters")
	if err != nil {
		return nil, err
	}
	return &AdodbStmt{c, s, val.ToIDispatch(), nil}, nil
}

func (s *AdodbStmt) Bind(bind []string) error {
	s.b = bind
	return nil
}

func (s *AdodbStmt) Close() error {
	s.s.Release()
	return nil
}

func (s *AdodbStmt) NumInput() int {
	if s.b != nil {
		return len(s.b)
	}
	_, err := oleutil.CallMethod(s.ps, "Refresh")
	if err != nil {
		return -1
	}
	val, err := oleutil.GetProperty(s.ps, "Count")
	if err != nil {
		return -1
	}
	c := int(val.Val)
	return c
}

func (s *AdodbStmt) bind(args []driver.Value) error {
	if s.b != nil {
		for i, v := range args {
			var b string = "?"
			if len(s.b) < i {
				b = s.b[i]
			}
			unknown, err := oleutil.CallMethod(s.s, "CreateParameter", b, 12, 1)
			if err != nil {
				return err
			}
			param := unknown.ToIDispatch()
			defer param.Release()
			_, err = oleutil.PutProperty(param, "Value", v)
			if err != nil {
				return err
			}
			_, err = oleutil.CallMethod(s.ps, "Append", param)
			if err != nil {
				return err
			}
		}
	} else {
		for i, v := range args {
			var varval ole.VARIANT
			varval.VT = ole.VT_I4
			varval.Val = int64(i)
			val, err := oleutil.CallMethod(s.ps, "Item", &varval)
			if err != nil {
				return err
			}
			item := val.ToIDispatch()
			defer item.Release()
			_, err = oleutil.PutProperty(item, "Value", v)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func (s *AdodbStmt) Query(args []driver.Value) (driver.Rows, error) {
	if err := s.bind(args); err != nil {
		return nil, err
	}
	rc, err := oleutil.CallMethod(s.s, "Execute")
	if err != nil {
		return nil, err
	}
	return &AdodbRows{s, rc.ToIDispatch(), -1, nil}, nil
}

func (s *AdodbStmt) Exec(args []driver.Value) (driver.Result, error) {
	if err := s.bind(args); err != nil {
		return nil, err
	}
	_, err := oleutil.CallMethod(s.s, "Execute")
	if err != nil {
		return nil, err
	}
	return driver.ResultNoRows, nil
}

type AdodbRows struct {
	s    *AdodbStmt
	rc   *ole.IDispatch
	nc   int
	cols []string
}

func (rc *AdodbRows) Close() error {
	_, err := oleutil.CallMethod(rc.rc, "Close")
	if err != nil {
		return err
	}
	return nil
}

func (rc *AdodbRows) Columns() []string {
	if rc.nc != len(rc.cols) {
		unknown, err := oleutil.GetProperty(rc.rc, "Fields")
		if err != nil {
			return []string{}
		}
		fields := unknown.ToIDispatch()
		defer fields.Release()
		val, err := oleutil.GetProperty(fields, "Count")
		if err != nil {
			return []string{}
		}
		rc.nc = int(val.Val)
		rc.cols = make([]string, rc.nc)
		for i := 0; i < rc.nc; i++ {
			var varval ole.VARIANT
			varval.VT = ole.VT_I4
			varval.Val = int64(i)
			val, err := oleutil.CallMethod(fields, "Item", &varval)
			if err != nil {
				return []string{}
			}
			item := val.ToIDispatch()
			if err != nil {
				return []string{}
			}
			name, err := oleutil.GetProperty(item, "Name")
			if err != nil {
				return []string{}
			}
			rc.cols[i] = name.ToString()
			item.Release()
		}
	}
	return rc.cols
}

func (rc *AdodbRows) Next(dest []driver.Value) error {
	unknown, err := oleutil.GetProperty(rc.rc, "EOF")
	if err != nil {
		return err
	}
	if unknown.Val != 0 {
		return errors.New("EOF")
	}
	_, err = oleutil.CallMethod(rc.rc, "MoveNext")
	if err != nil {
		fmt.Println(err)
		return err
	}
	unknown, err = oleutil.GetProperty(rc.rc, "Fields")
	if err != nil {
		return err
	}
	fields := unknown.ToIDispatch()
	defer fields.Release()
	for i := range dest {
		var varval ole.VARIANT
		varval.VT = ole.VT_I4
		varval.Val = int64(i)
		val, err := oleutil.CallMethod(fields, "Item", &varval)
		if err != nil {
			return err
		}
		field := val.ToIDispatch()
		defer field.Release()
		typ, err := oleutil.GetProperty(field, "Type")
		if err != nil {
			return err
		}
		val, err = oleutil.GetProperty(field, "Value")
		if err != nil {
			return err
		}
		field.Release()
		switch typ.Val {
		case 0: // ADEMPTY
			// TODO
		case 2: // ADSMALLINT
			dest[i] = int16(val.Val)
		case 3: // ADINTEGER
			dest[i] = int32(val.Val)
		case 4: // ADSINGLE
			dest[i] = float32(val.Val)
		case 5: // ADDOUBLE
			dest[i] = val.Val
		case 6: // ADCURRENCY
			// TODO
		case 7: // ADDATE
			// TODO
		case 8: // ADBSTR
			dest[i] = val.ToString()
		case 9: // ADIDISPATCH
			dest[i] = val.ToIDispatch()
		case 10: // ADERROR
			// TODO
		case 11: // ADBOOLEAN
			if val.Val != 0 {
				dest[i] = true
			} else {
				dest[i] = false
			}
		case 12: // ADVARIANT
			dest[i] = val
		case 13: // ADIUNKNOWN
			dest[i] = val.ToIUnknown()
		case 14: // ADDECIMAL
			// TODO
		case 16: // ADTINYINT
			dest[i] = int8(val.Val)
		case 17: // ADUNSIGNEDTINYINT
			dest[i] = uint8(val.Val)
		case 18: // ADUNSIGNEDSMALLINT
			dest[i] = uint16(val.Val)
		case 19: // ADUNSIGNEDINT
			dest[i] = uint32(val.Val)
		case 20: // ADBIGINT
			dest[i] = big.NewInt(val.Val)
		case 21: // ADUNSIGNEDBIGINT
			// TODO
		case 72: // ADGUID
			// TODO
		case 128: // ADBINARY
			sa := *(**ole.SAFEARRAY)(unsafe.Pointer(&val.Val))
			dest[i] = (*[1 << 30]byte)(unsafe.Pointer(uintptr(sa.PvData)))[0:sa.CbElements]
		case 129: // ADCHAR
			dest[i] = uint8(val.Val)
		case 130: // ADWCHAR
			dest[i] = uint16(val.Val)
		case 131: // ADNUMERIC
			dest[i] = val.Val
		case 132: // ADUSERDEFINED
			dest[i] = uintptr(val.Val)
		case 133: // ADDBDATE
			dest[i] = time.Unix(0, val.Val).UTC()
		case 134: // ADDBTIME
			dest[i] = time.Unix(0, val.Val).UTC()
		case 135: // ADDBTIMESTAMP
			dest[i] = time.Unix(0, val.Val).UTC()
		case 136: // ADCHAPTER
			dest[i] = val.ToString()
		case 200: // ADVARCHAR
			dest[i] = val.ToString()
		case 201: // ADLONGVARCHAR
			dest[i] = val.ToString()
		case 202: // ADVARWCHAR
			dest[i] = val.ToString()
		case 203: // ADLONGVARWCHAR
			dest[i] = val.ToString()
		case 204: // ADVARBINARY
			// TODO
		case 205: // ADLONGVARBINARY
			// TODO
		}
	}
	return nil
}
<|endoftext|>


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

package nano

import (
	"errors"
	"fmt"
	"log"
	"net"
	"sync/atomic"
	"time"

	"github.com/lonnng/nano/internal/codec"
	"github.com/lonnng/nano/internal/message"
	"github.com/lonnng/nano/internal/packet"
	"github.com/lonnng/nano/session"
)

const agentWriteBacklog = 16

var (
	ErrBrokenPipe   = errors.New("broken low-level pipe")
	ErrBufferExceed = errors.New("session send buffer exceed")
)

// Agent corresponding a user, used for store raw conn information
type (
	agent struct {
		session *session.Session    // session
		conn    net.Conn            // low-level conn fd
		state   int32               // current agent state
		chDie   chan struct{}       // wait for close
		chSend  chan pendingMessage // push message queue
		lastAt  int64               // last heartbeat unix time stamp
		decoder *codec.Decoder      // binary decoder
	}

	pendingMessage struct {
		typ     message.MessageType // message type
		route   string              // message route(push)
		mid     uint                // response message id(response)
		payload interface{}         // payload
	}
)

// Create new agent instance
func newAgent(conn net.Conn) *agent {
	a := &agent{
		conn:    conn,
		state:   statusStart,
		chDie:   make(chan struct{}),
		lastAt:  time.Now().Unix(),
		chSend:  make(chan pendingMessage, agentWriteBacklog),
		decoder: codec.NewDecoder(),
	}

	// binding session
	s := session.New(a)
	a.session = s

	return a
}

// Push, implementation for session.NetworkEntity interface
func (a *agent) Push(route string, v interface{}) error {
	if a.status() == statusClosed {
		return ErrBrokenPipe
	}

	if len(a.chSend) >= agentWriteBacklog {
		return ErrBufferExceed
	}

	if env.debug {
		log.Println(fmt.Sprintf("Type=Push, UID=%d, Route=%s, Data=%+v", a.session.Uid(), route, v))
	}

	a.chSend <- pendingMessage{typ: message.Push, route: route, payload: v}
	return nil
}

// Response, implementation for session.NetworkEntity interface
// Response message to session
func (a *agent) Response(v interface{}) error {
	if a.status() == statusClosed {
		return ErrBrokenPipe
	}

	mid := a.session.LastRID
	if mid <= 0 {
		return ErrSessionOnNotify
	}

	if len(a.chSend) >= agentWriteBacklog {
		return ErrBufferExceed
	}

	if env.debug {
		log.Println(fmt.Sprintf("Type=Response, UID=%d, MID=%d, Data=%+v", a.session.Uid(), mid, v))
	}

	a.chSend <- pendingMessage{typ: message.Response, mid: mid, payload: v}
	return nil
}

// Close, implementation for session.NetworkEntity interface
// Close closes the agent, clean inner state and close low-level connection.
// Any blocked Read or Write operations will be unblocked and return errors.
func (a *agent) Close() error {
	if a.status() == statusClosed {
		return ErrClosedSession
	}
	a.setStatus(statusClosed)

	if env.debug {
		log.Println(fmt.Sprintf("Session closed, Id=%d, IP=%s", a.session.ID(), a.conn.RemoteAddr()))
	}

	// close all channel
	close(a.chDie)
	return a.conn.Close()
}

// RemoteAddr, implementation for session.NetworkEntity interface
// returns the remote network address.
func (a *agent) RemoteAddr() net.Addr {
	return a.conn.RemoteAddr()
}

// String, implementation for Stringer interface
func (a *agent) String() string {
	return fmt.Sprintf("Remote=%s, LastTime=%d", a.conn.RemoteAddr().String(), a.lastAt)
}

func (a *agent) status() int32 {
	return atomic.LoadInt32(&a.state)
}

func (a *agent) setStatus(state int32) {
	atomic.StoreInt32(&a.state, state)
}

func (a *agent) write() {
	ticker := time.NewTicker(env.heartbeat)
	chWrite := make(chan []byte, agentWriteBacklog)
	// clean func
	defer func() {
		ticker.Stop()
		close(a.chSend)
		close(chWrite)
		a.Close()
		if env.debug {
			log.Println(fmt.Sprintf("Session write goroutine exit, SessionID=%d, UID=%d", a.session.ID(), a.session.Uid()))
		}
	}()

	for {
		select {
		case <-ticker.C:
			deadline := time.Now().Add(-2 * env.heartbeat).Unix()
			if a.lastAt < deadline {
				log.Println(fmt.Sprintf("Session heartbeat timeout, LastTime=%d, Deadline=%d", a.lastAt, deadline))
				return
			}
			chWrite <- hbd

		case data := <-chWrite:
			// close agent while low-level conn broken
			if _, err := a.conn.Write(data); err != nil {
				log.Println(err.Error())
				return
			}

		case data := <-a.chSend:
			payload, err := serializeOrRaw(data.payload)
			if err != nil {
				log.Println(err.Error())
				break
			}

			// construct message and encode
			m := &message.Message{
				Type:  data.typ,
				Data:  payload,
				Route: data.route,
				ID:    data.mid,
			}
			em, err := m.Encode()
			if err != nil {
				log.Println(err.Error())
				break
			}

			// packet encode
			p, err := codec.Encode(packet.Data, em)
			if err != nil {
				log.Println(err)
				break
			}
			chWrite <- p

		case <-a.chDie: // agent closed signal
			return

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

package nano

import (
	"errors"
	"fmt"
	"log"
	"net"
	"sync/atomic"
	"time"

	"github.com/lonnng/nano/internal/codec"
	"github.com/lonnng/nano/internal/message"
	"github.com/lonnng/nano/internal/packet"
	"github.com/lonnng/nano/session"
)

const agentWriteBacklog = 16

var (
	ErrBrokenPipe   = errors.New("broken low-level pipe")
	ErrBufferExceed = errors.New("session send buffer exceed")
)

// Agent corresponding a user, used for store raw conn information
type (
	agent struct {
		session *session.Session    // session
		conn    net.Conn            // low-level conn fd
		state   int32               // current agent state
		chDie   chan struct{}       // wait for close
		chSend  chan pendingMessage // push message queue
		lastAt  int64               // last heartbeat unix time stamp
		decoder *codec.Decoder      // binary decoder
	}

	pendingMessage struct {
		typ     message.MessageType // message type
		route   string              // message route(push)
		mid     uint                // response message id(response)
		payload interface{}         // payload
	}
)

// Create new agent instance
func newAgent(conn net.Conn) *agent {
	a := &agent{
		conn:    conn,
		state:   statusStart,
		chDie:   make(chan struct{}),
		lastAt:  time.Now().Unix(),
		chSend:  make(chan pendingMessage, agentWriteBacklog),
		decoder: codec.NewDecoder(),
	}

	// binding session
	s := session.New(a)
	a.session = s

	return a
}

// Push, implementation for session.NetworkEntity interface
func (a *agent) Push(route string, v interface{}) error {
	if a.status() == statusClosed {
		return ErrBrokenPipe
	}

	if len(a.chSend) >= agentWriteBacklog {
		return ErrBufferExceed
	}

	if env.debug {
		log.Println(fmt.Sprintf("Type=Push, UID=%d, Route=%s, Data=%+v", a.session.Uid(), route, v))
	}

	a.chSend <- pendingMessage{typ: message.Push, route: route, payload: v}
	return nil
}

// Response, implementation for session.NetworkEntity interface
// Response message to session
func (a *agent) Response(v interface{}) error {
	if a.status() == statusClosed {
		return ErrBrokenPipe
	}

	mid := a.session.LastRID
	if mid <= 0 {
		return ErrSessionOnNotify
	}

	if len(a.chSend) >= agentWriteBacklog {
		return ErrBufferExceed
	}

	if env.debug {
		log.Println(fmt.Sprintf("Type=Response, UID=%d, MID=%d, Data=%+v", a.session.Uid(), mid, v))
	}

	a.chSend <- pendingMessage{typ: message.Response, mid: mid, payload: v}
	return nil
}

// Close, implementation for session.NetworkEntity interface
// Close closes the agent, clean inner state and close low-level connection.
// Any blocked Read or Write operations will be unblocked and return errors.
func (a *agent) Close() error {
	if a.status() == statusClosed {
		return ErrClosedSession
	}
	a.setStatus(statusClosed)

	if env.debug {
		log.Println(fmt.Sprintf("Session closed, Id=%d, IP=%s", a.session.ID(), a.conn.RemoteAddr()))
	}

	// close all channel
	close(a.chDie)
	handler.chCloseSession <- a.session

	return a.conn.Close()
}

// RemoteAddr, implementation for session.NetworkEntity interface
// returns the remote network address.
func (a *agent) RemoteAddr() net.Addr {
	return a.conn.RemoteAddr()
}

// String, implementation for Stringer interface
func (a *agent) String() string {
	return fmt.Sprintf("Remote=%s, LastTime=%d", a.conn.RemoteAddr().String(), a.lastAt)
}

func (a *agent) status() int32 {
	return atomic.LoadInt32(&a.state)
}

func (a *agent) setStatus(state int32) {
	atomic.StoreInt32(&a.state, state)
}

func (a *agent) write() {
	ticker := time.NewTicker(env.heartbeat)
	chWrite := make(chan []byte, agentWriteBacklog)
	// clean func
	defer func() {
		ticker.Stop()
		close(a.chSend)
		close(chWrite)
		a.Close()
		if env.debug {
			log.Println(fmt.Sprintf("Session write goroutine exit, SessionID=%d, UID=%d", a.session.ID(), a.session.Uid()))
		}
	}()

	for {
		select {
		case <-ticker.C:
			deadline := time.Now().Add(-2 * env.heartbeat).Unix()
			if a.lastAt < deadline {
				log.Println(fmt.Sprintf("Session heartbeat timeout, LastTime=%d, Deadline=%d", a.lastAt, deadline))
				return
			}
			chWrite <- hbd

		case data := <-chWrite:
			// close agent while low-level conn broken
			if _, err := a.conn.Write(data); err != nil {
				log.Println(err.Error())
				return
			}

		case data := <-a.chSend:
			payload, err := serializeOrRaw(data.payload)
			if err != nil {
				log.Println(err.Error())
				break
			}

			// construct message and encode
			m := &message.Message{
				Type:  data.typ,
				Data:  payload,
				Route: data.route,
				ID:    data.mid,
			}
			em, err := m.Encode()
			if err != nil {
				log.Println(err.Error())
				break
			}

			// packet encode
			p, err := codec.Encode(packet.Data, em)
			if err != nil {
				log.Println(err)
				break
			}
			chWrite <- p

		case <-a.chDie: // agent closed signal
			return

		case <-env.die: // application quit
			return
		}
	}
}
<|endoftext|>


<commit_before>package telegraf

import (
	"fmt"
	"log"
	"os"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/influxdb/telegraf/outputs"
	"github.com/influxdb/telegraf/plugins"
)

type runningOutput struct {
	name   string
	output outputs.Output
}

type runningPlugin struct {
	name   string
	plugin plugins.Plugin
	config *ConfiguredPlugin
}

// Agent runs telegraf and collects data based on the given config
type Agent struct {

	// Interval at which to gather information
	Interval Duration

	// Run in debug mode?
	Debug    bool
	Hostname string

	Config *Config

	outputs []*runningOutput
	plugins []*runningPlugin
}

// NewAgent returns an Agent struct based off the given Config
func NewAgent(config *Config) (*Agent, error) {
	agent := &Agent{Config: config, Interval: Duration{10 * time.Second}}

	err := config.ApplyAgent(agent)
	if err != nil {
		return nil, err
	}

	if agent.Hostname == "" {
		hostname, err := os.Hostname()
		if err != nil {
			return nil, err
		}

		agent.Hostname = hostname
	}

	if config.Tags == nil {
		config.Tags = map[string]string{}
	}

	config.Tags["host"] = agent.Hostname

	return agent, nil
}

// Connect connects to all configured outputs
func (a *Agent) Connect() error {
	for _, o := range a.outputs {
		err := o.output.Connect()
		if err != nil {
			return err
		}
	}
	return nil
}

// Close closes the connection to all configured outputs
func (a *Agent) Close() error {
	var err error
	for _, o := range a.outputs {
		err = o.output.Close()
	}
	return err
}

// LoadOutputs loads the agent's outputs
func (a *Agent) LoadOutputs() ([]string, error) {
	var names []string

	for _, name := range a.Config.OutputsDeclared() {
		creator, ok := outputs.Outputs[name]
		if !ok {
			return nil, fmt.Errorf("Undefined but requested output: %s", name)
		}

		output := creator()

		err := a.Config.ApplyOutput(name, output)
		if err != nil {
			return nil, err
		}

		a.outputs = append(a.outputs, &runningOutput{name, output})
		names = append(names, name)
	}

	sort.Strings(names)

	return names, nil
}

// LoadPlugins loads the agent's plugins
func (a *Agent) LoadPlugins(pluginsFilter string) ([]string, error) {
	var names []string
	var filters []string

	pluginsFilter = strings.TrimSpace(pluginsFilter)
	if pluginsFilter != "" {
		filters = strings.Split(":"+pluginsFilter+":", ":")
	}

	for _, name := range a.Config.PluginsDeclared() {
		creator, ok := plugins.Plugins[name]
		if !ok {
			return nil, fmt.Errorf("Undefined but requested plugin: %s", name)
		}

		isPluginEnabled := false
		if len(filters) > 0 {
			for _, runeValue := range filters {
				if runeValue != "" && strings.ToLower(runeValue) == strings.ToLower(name) {
					fmt.Printf("plugin [%s] is enabled (filter options)\n", name)
					isPluginEnabled = true
					break
				}
			}
		} else {
			// if no filter, we ALWAYS accept the plugin
			isPluginEnabled = true
		}

		if isPluginEnabled {
			plugin := creator()
			config, err := a.Config.ApplyPlugin(name, plugin)
			if err != nil {
				return nil, err
			}

			a.plugins = append(a.plugins, &runningPlugin{name, plugin, config})
			names = append(names, name)
		}
	}

	sort.Strings(names)

	return names, nil
}

func (a *Agent) crankParallel() error {
	points := make(chan *BatchPoints, len(a.plugins))

	var wg sync.WaitGroup

	for _, plugin := range a.plugins {
		if plugin.config.Interval != 0 {
			continue
		}

		wg.Add(1)
		go func(plugin *runningPlugin) {
			defer wg.Done()

			var acc BatchPoints
			acc.Debug = a.Debug
			acc.Prefix = plugin.name + "_"
			acc.Config = plugin.config

			plugin.plugin.Gather(&acc)

			points <- &acc
		}(plugin)
	}

	wg.Wait()

	close(points)

	var bp BatchPoints
	bp.Time = time.Now()
	bp.Tags = a.Config.Tags

	for sub := range points {
		bp.Points = append(bp.Points, sub.Points...)
	}

	return a.flush(&bp)
}

func (a *Agent) crank() error {
	var bp BatchPoints

	bp.Debug = a.Debug

	for _, plugin := range a.plugins {
		bp.Prefix = plugin.name + "_"
		bp.Config = plugin.config
		err := plugin.plugin.Gather(&bp)
		if err != nil {
			return err
		}
	}

	bp.Time = time.Now()
	bp.Tags = a.Config.Tags

	return a.flush(&bp)
}

func (a *Agent) crankSeparate(shutdown chan struct{}, plugin *runningPlugin) error {
	ticker := time.NewTicker(plugin.config.Interval)

	for {
		var bp BatchPoints

		bp.Debug = a.Debug

		bp.Prefix = plugin.name + "_"
		bp.Config = plugin.config
		err := plugin.plugin.Gather(&bp)
		if err != nil {
			return err
		}

		bp.Tags = a.Config.Tags
		bp.Time = time.Now()

		err = a.flush(&bp)
		if err != nil {
			return err
		}

		select {
		case <-shutdown:
			return nil
		case <-ticker.C:
			continue
		}
	}
}

func (a *Agent) flush(bp *BatchPoints) error {
	var wg sync.WaitGroup
	var outerr error
	for _, o := range a.outputs {
		wg.Add(1)
		go func(ro *runningOutput) {
			defer wg.Done()
			outerr = ro.output.Write(bp.BatchPoints)
		}(o)
	}

	wg.Wait()

	return outerr
}

// TestAllPlugins verifies that we can 'Gather' from all plugins with the
// default configuration
func (a *Agent) TestAllPlugins() error {
	var names []string

	for name := range plugins.Plugins {
		names = append(names, name)
	}

	sort.Strings(names)

	var acc BatchPoints
	acc.Debug = true

	fmt.Printf("* Testing all plugins with default configuration\n")

	for _, name := range names {
		plugin := plugins.Plugins[name]()

		fmt.Printf("* Plugin: %s\n", name)

		acc.Prefix = name + "_"
		err := plugin.Gather(&acc)
		if err != nil {
			return err
		}
	}

	return nil
}

// Test verifies that we can 'Gather' from all plugins with their configured
// Config struct
func (a *Agent) Test() error {
	var acc BatchPoints

	acc.Debug = true

	for _, plugin := range a.plugins {
		acc.Prefix = plugin.name + "_"
		acc.Config = plugin.config

		fmt.Printf("* Plugin: %s\n", plugin.name)
		if plugin.config.Interval != 0 {
			fmt.Printf("* Internal: %s\n", plugin.config.Interval)
		}

		err := plugin.plugin.Gather(&acc)
		if err != nil {
			return err
		}
	}

	return nil
}

// Run runs the agent daemon, gathering every Interval
func (a *Agent) Run(shutdown chan struct{}) error {
	var wg sync.WaitGroup

	for _, plugin := range a.plugins {
		if plugin.config.Interval != 0 {
			wg.Add(1)
			go func(plugin *runningPlugin) {
				defer wg.Done()
				a.crankSeparate(shutdown, plugin)
			}(plugin)
		}
	}

	defer wg.Wait()

	ticker := time.NewTicker(a.Interval.Duration)

	for {
		err := a.crankParallel()
		if err != nil {
			log.Printf("Error in plugins: %s", err)
		}

		select {
		case <-shutdown:
			return nil
		case <-ticker.C:
			continue
		}
	}
}
<commit_msg>Log plugin errors in crankParallel and crankSeparate cases. Previously errors weren't logged in these cases.<commit_after>package telegraf

import (
	"fmt"
	"log"
	"os"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/influxdb/telegraf/outputs"
	"github.com/influxdb/telegraf/plugins"
)

type runningOutput struct {
	name   string
	output outputs.Output
}

type runningPlugin struct {
	name   string
	plugin plugins.Plugin
	config *ConfiguredPlugin
}

// Agent runs telegraf and collects data based on the given config
type Agent struct {

	// Interval at which to gather information
	Interval Duration

	// Run in debug mode?
	Debug    bool
	Hostname string

	Config *Config

	outputs []*runningOutput
	plugins []*runningPlugin
}

// NewAgent returns an Agent struct based off the given Config
func NewAgent(config *Config) (*Agent, error) {
	agent := &Agent{Config: config, Interval: Duration{10 * time.Second}}

	err := config.ApplyAgent(agent)
	if err != nil {
		return nil, err
	}

	if agent.Hostname == "" {
		hostname, err := os.Hostname()
		if err != nil {
			return nil, err
		}

		agent.Hostname = hostname
	}

	if config.Tags == nil {
		config.Tags = map[string]string{}
	}

	config.Tags["host"] = agent.Hostname

	return agent, nil
}

// Connect connects to all configured outputs
func (a *Agent) Connect() error {
	for _, o := range a.outputs {
		err := o.output.Connect()
		if err != nil {
			return err
		}
	}
	return nil
}

// Close closes the connection to all configured outputs
func (a *Agent) Close() error {
	var err error
	for _, o := range a.outputs {
		err = o.output.Close()
	}
	return err
}

// LoadOutputs loads the agent's outputs
func (a *Agent) LoadOutputs() ([]string, error) {
	var names []string

	for _, name := range a.Config.OutputsDeclared() {
		creator, ok := outputs.Outputs[name]
		if !ok {
			return nil, fmt.Errorf("Undefined but requested output: %s", name)
		}

		output := creator()

		err := a.Config.ApplyOutput(name, output)
		if err != nil {
			return nil, err
		}

		a.outputs = append(a.outputs, &runningOutput{name, output})
		names = append(names, name)
	}

	sort.Strings(names)

	return names, nil
}

// LoadPlugins loads the agent's plugins
func (a *Agent) LoadPlugins(pluginsFilter string) ([]string, error) {
	var names []string
	var filters []string

	pluginsFilter = strings.TrimSpace(pluginsFilter)
	if pluginsFilter != "" {
		filters = strings.Split(":"+pluginsFilter+":", ":")
	}

	for _, name := range a.Config.PluginsDeclared() {
		creator, ok := plugins.Plugins[name]
		if !ok {
			return nil, fmt.Errorf("Undefined but requested plugin: %s", name)
		}

		isPluginEnabled := false
		if len(filters) > 0 {
			for _, runeValue := range filters {
				if runeValue != "" && strings.ToLower(runeValue) == strings.ToLower(name) {
					fmt.Printf("plugin [%s] is enabled (filter options)\n", name)
					isPluginEnabled = true
					break
				}
			}
		} else {
			// if no filter, we ALWAYS accept the plugin
			isPluginEnabled = true
		}

		if isPluginEnabled {
			plugin := creator()
			config, err := a.Config.ApplyPlugin(name, plugin)
			if err != nil {
				return nil, err
			}

			a.plugins = append(a.plugins, &runningPlugin{name, plugin, config})
			names = append(names, name)
		}
	}

	sort.Strings(names)

	return names, nil
}

func (a *Agent) crankParallel() error {
	points := make(chan *BatchPoints, len(a.plugins))

	var wg sync.WaitGroup

	for _, plugin := range a.plugins {
		if plugin.config.Interval != 0 {
			continue
		}

		wg.Add(1)
		go func(plugin *runningPlugin) {
			defer wg.Done()

			var acc BatchPoints
			acc.Debug = a.Debug
			acc.Prefix = plugin.name + "_"
			acc.Config = plugin.config

			err := plugin.plugin.Gather(&acc)
			if err != nil {
				log.Printf("Error in plugins: %s", err)
			}

			points <- &acc
		}(plugin)
	}

	wg.Wait()

	close(points)

	var bp BatchPoints
	bp.Time = time.Now()
	bp.Tags = a.Config.Tags

	for sub := range points {
		bp.Points = append(bp.Points, sub.Points...)
	}

	return a.flush(&bp)
}

func (a *Agent) crank() error {
	var bp BatchPoints

	bp.Debug = a.Debug

	for _, plugin := range a.plugins {
		bp.Prefix = plugin.name + "_"
		bp.Config = plugin.config
		err := plugin.plugin.Gather(&bp)
		if err != nil {
			return err
		}
	}

	bp.Time = time.Now()
	bp.Tags = a.Config.Tags

	return a.flush(&bp)
}

func (a *Agent) crankSeparate(shutdown chan struct{}, plugin *runningPlugin) error {
	ticker := time.NewTicker(plugin.config.Interval)

	for {
		var bp BatchPoints

		bp.Debug = a.Debug

		bp.Prefix = plugin.name + "_"
		bp.Config = plugin.config
		err := plugin.plugin.Gather(&bp)
		if err != nil {
			return err
		}

		bp.Tags = a.Config.Tags
		bp.Time = time.Now()

		err = a.flush(&bp)
		if err != nil {
			return err
		}

		select {
		case <-shutdown:
			return nil
		case <-ticker.C:
			continue
		}
	}
}

func (a *Agent) flush(bp *BatchPoints) error {
	var wg sync.WaitGroup
	var outerr error
	for _, o := range a.outputs {
		wg.Add(1)
		go func(ro *runningOutput) {
			defer wg.Done()
			outerr = ro.output.Write(bp.BatchPoints)
		}(o)
	}

	wg.Wait()

	return outerr
}

// TestAllPlugins verifies that we can 'Gather' from all plugins with the
// default configuration
func (a *Agent) TestAllPlugins() error {
	var names []string

	for name := range plugins.Plugins {
		names = append(names, name)
	}

	sort.Strings(names)

	var acc BatchPoints
	acc.Debug = true

	fmt.Printf("* Testing all plugins with default configuration\n")

	for _, name := range names {
		plugin := plugins.Plugins[name]()

		fmt.Printf("* Plugin: %s\n", name)

		acc.Prefix = name + "_"
		err := plugin.Gather(&acc)
		if err != nil {
			return err
		}
	}

	return nil
}

// Test verifies that we can 'Gather' from all plugins with their configured
// Config struct
func (a *Agent) Test() error {
	var acc BatchPoints

	acc.Debug = true

	for _, plugin := range a.plugins {
		acc.Prefix = plugin.name + "_"
		acc.Config = plugin.config

		fmt.Printf("* Plugin: %s\n", plugin.name)
		if plugin.config.Interval != 0 {
			fmt.Printf("* Internal: %s\n", plugin.config.Interval)
		}

		err := plugin.plugin.Gather(&acc)
		if err != nil {
			return err
		}
	}

	return nil
}

// Run runs the agent daemon, gathering every Interval
func (a *Agent) Run(shutdown chan struct{}) error {
	var wg sync.WaitGroup

	for _, plugin := range a.plugins {
		if plugin.config.Interval != 0 {
			wg.Add(1)
			go func(plugin *runningPlugin) {
				defer wg.Done()
				err := a.crankSeparate(shutdown, plugin)
				if err != nil {
					log.Printf("Error in plugins: %s", err)
				}
			}(plugin)
		}
	}

	defer wg.Wait()

	ticker := time.NewTicker(a.Interval.Duration)

	for {
		err := a.crankParallel()
		if err != nil {
			log.Printf("Error in plugins: %s", err)
		}

		select {
		case <-shutdown:
			return nil
		case <-ticker.C:
			continue
		}
	}
}
<|endoftext|>


<commit_before>package gosnmp

type Agent struct {
	snmpContext
}

func NewAgent(name string, maxTargets int, logger Logger) *Agent {
	return NewAgentWithPort(name, maxTargets, 161, logger)
}

func NewAgentWithPort(name string, maxTargets int, port int, logger Logger) *Agent {
	agent := new(Agent)
	agent.snmpContext = *newContext(name, maxTargets, false, port, logger)
	return agent
}

func (ctxt *snmpContext) processIncomingRequest(req SnmpRequest) {

}
<commit_msg>Add BasicOidHandler interface<commit_after>package gosnmp

type Agent struct {
	snmpContext
}

func NewAgent(name string, maxTargets int, logger Logger) *Agent {
	return NewAgentWithPort(name, maxTargets, 161, logger)
}

func NewAgentWithPort(name string, maxTargets int, port int, logger Logger) *Agent {
	agent := new(Agent)
	agent.snmpContext = *newContext(name, maxTargets, false, port, logger)
	return agent
}

func (ctxt *snmpContext) processIncomingRequest(req SnmpRequest) {

}

type BasicOidHandler interface {
	Get(Varbind) error
	Set(Varbind) error
}

func (agent *Agent) register<|endoftext|>


<commit_before>package fzf

import (
	"fmt"
	"runtime"
	"sort"
	"time"
)

type MatchRequest struct {
	chunks  []*Chunk
	pattern *Pattern
}

type Matcher struct {
	patternBuilder func([]rune) *Pattern
	sort           bool
	eventBox       *EventBox
	reqBox         *EventBox
	partitions     int
	queryCache     QueryCache
}

const (
	REQ_RETRY EventType = iota
	REQ_RESET
)

const (
	STAT_CANCELLED int = iota
	STAT_QCH
	STAT_CHUNKS
)

const (
	PROGRESS_MIN_DURATION = 200 * time.Millisecond
)

func NewMatcher(patternBuilder func([]rune) *Pattern,
	sort bool, eventBox *EventBox) *Matcher {
	return &Matcher{
		patternBuilder: patternBuilder,
		sort:           sort,
		eventBox:       eventBox,
		reqBox:         NewEventBox(),
		partitions:     runtime.NumCPU(),
		queryCache:     make(QueryCache)}
}

func (m *Matcher) Loop() {
	prevCount := 0

	for {
		var request MatchRequest

		m.reqBox.Wait(func(events *Events) {
			for _, val := range *events {
				switch val := val.(type) {
				case MatchRequest:
					request = val
				default:
					panic(fmt.Sprintf("Unexpected type: %T", val))
				}
			}
			events.Clear()
		})

		// Restart search
		patternString := request.pattern.AsString()
		allMatches := []*Item{}
		cancelled := false
		count := CountItems(request.chunks)

		foundCache := false
		if count == prevCount {
			// Look up queryCache
			if cached, found := m.queryCache[patternString]; found {
				foundCache = true
				allMatches = cached
			}
		} else {
			// Invalidate queryCache
			prevCount = count
			m.queryCache = make(QueryCache)
		}

		if !foundCache {
			allMatches, cancelled = m.scan(request, 0)
		}

		if !cancelled {
			m.queryCache[patternString] = allMatches
			m.eventBox.Set(EVT_SEARCH_FIN, allMatches)
		}
	}
}

func (m *Matcher) sliceChunks(chunks []*Chunk) [][]*Chunk {
	perSlice := len(chunks) / m.partitions

	// No need to parallelize
	if perSlice == 0 {
		return [][]*Chunk{chunks}
	}

	slices := make([][]*Chunk, m.partitions)
	for i := 0; i < m.partitions; i++ {
		start := i * perSlice
		end := start + perSlice
		if i == m.partitions-1 {
			end = len(chunks)
		}
		slices[i] = chunks[start:end]
	}
	return slices
}

type partialResult struct {
	index   int
	matches []*Item
}

func (m *Matcher) scan(request MatchRequest, limit int) ([]*Item, bool) {
	startedAt := time.Now()

	numChunks := len(request.chunks)
	if numChunks == 0 {
		return []*Item{}, false
	}
	pattern := request.pattern
	empty := pattern.IsEmpty()
	cancelled := NewAtomicBool(false)

	slices := m.sliceChunks(request.chunks)
	numSlices := len(slices)
	resultChan := make(chan partialResult, numSlices)
	countChan := make(chan int, numSlices)

	for idx, chunks := range slices {
		go func(idx int, chunks []*Chunk) {
			sliceMatches := []*Item{}
			for _, chunk := range chunks {
				var matches []*Item
				if empty {
					matches = *chunk
				} else {
					matches = request.pattern.Match(chunk)
				}
				sliceMatches = append(sliceMatches, matches...)
				if cancelled.Get() {
					return
				}
				countChan <- len(sliceMatches)
			}
			if !empty && m.sort {
				sort.Sort(ByRelevance(sliceMatches))
			}
			resultChan <- partialResult{idx, sliceMatches}
		}(idx, chunks)
	}

	count := 0
	matchCount := 0
	for matchesInChunk := range countChan {
		count += 1
		matchCount += matchesInChunk

		if limit > 0 && matchCount > limit {
			return nil, true // For --select-1 and --exit-0
		}

		if count == numChunks {
			break
		}

		if !empty && m.reqBox.Peak(REQ_RESET) {
			cancelled.Set(true)
			return nil, true
		}

		if time.Now().Sub(startedAt) > PROGRESS_MIN_DURATION {
			m.eventBox.Set(EVT_SEARCH_PROGRESS, float32(count)/float32(numChunks))
		}
	}

	partialResults := make([][]*Item, numSlices)
	for range slices {
		partialResult := <-resultChan
		partialResults[partialResult.index] = partialResult.matches
	}

	var allMatches []*Item
	if empty || !m.sort {
		allMatches = []*Item{}
		for _, matches := range partialResults {
			allMatches = append(allMatches, matches...)
		}
	} else {
		allMatches = SortMerge(partialResults)
	}

	return allMatches, false
}

func (m *Matcher) Reset(chunks []*Chunk, patternRunes []rune, cancel bool) {
	pattern := m.patternBuilder(patternRunes)

	var event EventType
	if cancel {
		event = REQ_RESET
	} else {
		event = REQ_RETRY
	}
	m.reqBox.Set(event, MatchRequest{chunks, pattern})
}
<commit_msg>Fix scan limit for --select-1 and --exit-0 options<commit_after>package fzf

import (
	"fmt"
	"runtime"
	"sort"
	"time"
)

type MatchRequest struct {
	chunks  []*Chunk
	pattern *Pattern
}

type Matcher struct {
	patternBuilder func([]rune) *Pattern
	sort           bool
	eventBox       *EventBox
	reqBox         *EventBox
	partitions     int
	queryCache     QueryCache
}

const (
	REQ_RETRY EventType = iota
	REQ_RESET
)

const (
	STAT_CANCELLED int = iota
	STAT_QCH
	STAT_CHUNKS
)

const (
	PROGRESS_MIN_DURATION = 200 * time.Millisecond
)

func NewMatcher(patternBuilder func([]rune) *Pattern,
	sort bool, eventBox *EventBox) *Matcher {
	return &Matcher{
		patternBuilder: patternBuilder,
		sort:           sort,
		eventBox:       eventBox,
		reqBox:         NewEventBox(),
		partitions:     runtime.NumCPU(),
		queryCache:     make(QueryCache)}
}

func (m *Matcher) Loop() {
	prevCount := 0

	for {
		var request MatchRequest

		m.reqBox.Wait(func(events *Events) {
			for _, val := range *events {
				switch val := val.(type) {
				case MatchRequest:
					request = val
				default:
					panic(fmt.Sprintf("Unexpected type: %T", val))
				}
			}
			events.Clear()
		})

		// Restart search
		patternString := request.pattern.AsString()
		allMatches := []*Item{}
		cancelled := false
		count := CountItems(request.chunks)

		foundCache := false
		if count == prevCount {
			// Look up queryCache
			if cached, found := m.queryCache[patternString]; found {
				foundCache = true
				allMatches = cached
			}
		} else {
			// Invalidate queryCache
			prevCount = count
			m.queryCache = make(QueryCache)
		}

		if !foundCache {
			allMatches, cancelled = m.scan(request, 0)
		}

		if !cancelled {
			m.queryCache[patternString] = allMatches
			m.eventBox.Set(EVT_SEARCH_FIN, allMatches)
		}
	}
}

func (m *Matcher) sliceChunks(chunks []*Chunk) [][]*Chunk {
	perSlice := len(chunks) / m.partitions

	// No need to parallelize
	if perSlice == 0 {
		return [][]*Chunk{chunks}
	}

	slices := make([][]*Chunk, m.partitions)
	for i := 0; i < m.partitions; i++ {
		start := i * perSlice
		end := start + perSlice
		if i == m.partitions-1 {
			end = len(chunks)
		}
		slices[i] = chunks[start:end]
	}
	return slices
}

type partialResult struct {
	index   int
	matches []*Item
}

func (m *Matcher) scan(request MatchRequest, limit int) ([]*Item, bool) {
	startedAt := time.Now()

	numChunks := len(request.chunks)
	if numChunks == 0 {
		return []*Item{}, false
	}
	pattern := request.pattern
	empty := pattern.IsEmpty()
	cancelled := NewAtomicBool(false)

	slices := m.sliceChunks(request.chunks)
	numSlices := len(slices)
	resultChan := make(chan partialResult, numSlices)
	countChan := make(chan int, numSlices)

	for idx, chunks := range slices {
		go func(idx int, chunks []*Chunk) {
			sliceMatches := []*Item{}
			for _, chunk := range chunks {
				var matches []*Item
				if empty {
					matches = *chunk
				} else {
					matches = request.pattern.Match(chunk)
				}
				sliceMatches = append(sliceMatches, matches...)
				if cancelled.Get() {
					return
				}
				countChan <- len(matches)
			}
			if !empty && m.sort {
				sort.Sort(ByRelevance(sliceMatches))
			}
			resultChan <- partialResult{idx, sliceMatches}
		}(idx, chunks)
	}

	count := 0
	matchCount := 0
	for matchesInChunk := range countChan {
		count += 1
		matchCount += matchesInChunk

		if limit > 0 && matchCount > limit {
			return nil, true // For --select-1 and --exit-0
		}

		if count == numChunks {
			break
		}

		if !empty && m.reqBox.Peak(REQ_RESET) {
			cancelled.Set(true)
			return nil, true
		}

		if time.Now().Sub(startedAt) > PROGRESS_MIN_DURATION {
			m.eventBox.Set(EVT_SEARCH_PROGRESS, float32(count)/float32(numChunks))
		}
	}

	partialResults := make([][]*Item, numSlices)
	for range slices {
		partialResult := <-resultChan
		partialResults[partialResult.index] = partialResult.matches
	}

	var allMatches []*Item
	if empty || !m.sort {
		allMatches = []*Item{}
		for _, matches := range partialResults {
			allMatches = append(allMatches, matches...)
		}
	} else {
		allMatches = SortMerge(partialResults)
	}

	return allMatches, false
}

func (m *Matcher) Reset(chunks []*Chunk, patternRunes []rune, cancel bool) {
	pattern := m.patternBuilder(patternRunes)

	var event EventType
	if cancel {
		event = REQ_RESET
	} else {
		event = REQ_RETRY
	}
	m.reqBox.Set(event, MatchRequest{chunks, pattern})
}
<|endoftext|>


<commit_before>package local

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"io"

	"github.com/docker/docker/api/types/plugins/logdriver"
	"github.com/docker/docker/daemon/logger"
	"github.com/docker/docker/daemon/logger/loggerutils"
	"github.com/docker/docker/errdefs"
	"github.com/pkg/errors"
)

// maxMsgLen is the maximum size of the logger.Message after serialization.
// logger.defaultBufSize caps the size of Line field.
const maxMsgLen int = 1e6 // 1MB.

func (d *driver) ReadLogs(config logger.ReadConfig) *logger.LogWatcher {
	return d.logfile.ReadLogs(config)
}

func getTailReader(ctx context.Context, r loggerutils.SizeReaderAt, req int) (io.Reader, int, error) {
	size := r.Size()
	if req < 0 {
		return nil, 0, errdefs.InvalidParameter(errors.Errorf("invalid number of lines to tail: %d", req))
	}

	if size < (encodeBinaryLen*2)+1 {
		return bytes.NewReader(nil), 0, nil
	}

	const encodeBinaryLen64 = int64(encodeBinaryLen)
	var found int

	buf := make([]byte, encodeBinaryLen)

	offset := size
	for {
		select {
		case <-ctx.Done():
			return nil, 0, ctx.Err()
		default:
		}

		n, err := r.ReadAt(buf, offset-encodeBinaryLen64)
		if err != nil && err != io.EOF {
			return nil, 0, errors.Wrap(err, "error reading log message footer")
		}

		if n != encodeBinaryLen {
			return nil, 0, errdefs.DataLoss(errors.New("unexpected number of bytes read from log message footer"))
		}

		msgLen := binary.BigEndian.Uint32(buf)

		n, err = r.ReadAt(buf, offset-encodeBinaryLen64-encodeBinaryLen64-int64(msgLen))
		if err != nil && err != io.EOF {
			return nil, 0, errors.Wrap(err, "error reading log message header")
		}

		if n != encodeBinaryLen {
			return nil, 0, errdefs.DataLoss(errors.New("unexpected number of bytes read from log message header"))
		}

		if msgLen != binary.BigEndian.Uint32(buf) {
			return nil, 0, errdefs.DataLoss(errors.Wrap(err, "log message header and footer indicate different message sizes"))
		}

		found++
		offset -= int64(msgLen)
		offset -= encodeBinaryLen64 * 2
		if found == req {
			break
		}
		if offset <= 0 {
			break
		}
	}

	return io.NewSectionReader(r, offset, size), found, nil
}

type decoder struct {
	rdr   io.Reader
	proto *logdriver.LogEntry
	// buf keeps bytes from rdr.
	buf []byte
	// offset is the position in buf.
	// If offset > 0, buf[offset:] has bytes which are read but haven't used.
	offset int
	// nextMsgLen is the length of the next log message.
	// If nextMsgLen = 0, a new value must be read from rdr.
	nextMsgLen int
}

func (d *decoder) readRecord(size int) error {
	var err error
	for i := 0; i < maxDecodeRetry; i++ {
		var n int
		n, err = io.ReadFull(d.rdr, d.buf[d.offset:size])
		d.offset += n
		if err != nil {
			if err != io.ErrUnexpectedEOF {
				return err
			}
			continue
		}
		break
	}
	if err != nil {
		return err
	}
	d.offset = 0
	return nil
}

func (d *decoder) Decode() (*logger.Message, error) {
	if d.proto == nil {
		d.proto = &logdriver.LogEntry{}
	} else {
		resetProto(d.proto)
	}
	if d.buf == nil {
		d.buf = make([]byte, initialBufSize)
	}

	if d.nextMsgLen == 0 {
		msgLen, err := d.decodeSizeHeader()
		if err != nil {
			return nil, err
		}

		if msgLen > maxMsgLen {
			return nil, fmt.Errorf("log message is too large (%d > %d)", msgLen, maxMsgLen)
		}

		if len(d.buf) < msgLen+encodeBinaryLen {
			d.buf = make([]byte, msgLen+encodeBinaryLen)
		} else if msgLen <= initialBufSize {
			d.buf = d.buf[:initialBufSize]
		} else {
			d.buf = d.buf[:msgLen+encodeBinaryLen]
		}

		d.nextMsgLen = msgLen
	}
	return d.decodeLogEntry()
}

func (d *decoder) Reset(rdr io.Reader) {
	if d.rdr == rdr {
		return
	}

	d.rdr = rdr
	if d.proto != nil {
		resetProto(d.proto)
	}
	if d.buf != nil {
		d.buf = d.buf[:initialBufSize]
	}
	d.offset = 0
	d.nextMsgLen = 0
}

func (d *decoder) Close() {
	d.buf = d.buf[:0]
	d.buf = nil
	if d.proto != nil {
		resetProto(d.proto)
	}
	d.rdr = nil
}

func decodeFunc(rdr io.Reader) loggerutils.Decoder {
	return &decoder{rdr: rdr}
}

func (d *decoder) decodeSizeHeader() (int, error) {
	err := d.readRecord(encodeBinaryLen)
	if err != nil {
		return 0, errors.Wrap(err, "could not read a size header")
	}

	msgLen := int(binary.BigEndian.Uint32(d.buf[:encodeBinaryLen]))
	return msgLen, nil
}

func (d *decoder) decodeLogEntry() (*logger.Message, error) {
	msgLen := d.nextMsgLen
	err := d.readRecord(msgLen + encodeBinaryLen)
	if err != nil {
		return nil, errors.Wrapf(err, "could not read a log entry (size=%d+%d)", msgLen, encodeBinaryLen)
	}
	d.nextMsgLen = 0

	if err := d.proto.Unmarshal(d.buf[:msgLen]); err != nil {
		return nil, errors.Wrapf(err, "error unmarshalling log entry (size=%d)", msgLen)
	}

	msg := protoToMessage(d.proto)
	if msg.PLogMetaData == nil {
		msg.Line = append(msg.Line, '\n')
	}

	return msg, nil
}
<commit_msg>daemon/logger/local: fix appending newlines<commit_after>package local

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"io"

	"github.com/docker/docker/api/types/plugins/logdriver"
	"github.com/docker/docker/daemon/logger"
	"github.com/docker/docker/daemon/logger/loggerutils"
	"github.com/docker/docker/errdefs"
	"github.com/pkg/errors"
)

// maxMsgLen is the maximum size of the logger.Message after serialization.
// logger.defaultBufSize caps the size of Line field.
const maxMsgLen int = 1e6 // 1MB.

func (d *driver) ReadLogs(config logger.ReadConfig) *logger.LogWatcher {
	return d.logfile.ReadLogs(config)
}

func getTailReader(ctx context.Context, r loggerutils.SizeReaderAt, req int) (io.Reader, int, error) {
	size := r.Size()
	if req < 0 {
		return nil, 0, errdefs.InvalidParameter(errors.Errorf("invalid number of lines to tail: %d", req))
	}

	if size < (encodeBinaryLen*2)+1 {
		return bytes.NewReader(nil), 0, nil
	}

	const encodeBinaryLen64 = int64(encodeBinaryLen)
	var found int

	buf := make([]byte, encodeBinaryLen)

	offset := size
	for {
		select {
		case <-ctx.Done():
			return nil, 0, ctx.Err()
		default:
		}

		n, err := r.ReadAt(buf, offset-encodeBinaryLen64)
		if err != nil && err != io.EOF {
			return nil, 0, errors.Wrap(err, "error reading log message footer")
		}

		if n != encodeBinaryLen {
			return nil, 0, errdefs.DataLoss(errors.New("unexpected number of bytes read from log message footer"))
		}

		msgLen := binary.BigEndian.Uint32(buf)

		n, err = r.ReadAt(buf, offset-encodeBinaryLen64-encodeBinaryLen64-int64(msgLen))
		if err != nil && err != io.EOF {
			return nil, 0, errors.Wrap(err, "error reading log message header")
		}

		if n != encodeBinaryLen {
			return nil, 0, errdefs.DataLoss(errors.New("unexpected number of bytes read from log message header"))
		}

		if msgLen != binary.BigEndian.Uint32(buf) {
			return nil, 0, errdefs.DataLoss(errors.Wrap(err, "log message header and footer indicate different message sizes"))
		}

		found++
		offset -= int64(msgLen)
		offset -= encodeBinaryLen64 * 2
		if found == req {
			break
		}
		if offset <= 0 {
			break
		}
	}

	return io.NewSectionReader(r, offset, size), found, nil
}

type decoder struct {
	rdr   io.Reader
	proto *logdriver.LogEntry
	// buf keeps bytes from rdr.
	buf []byte
	// offset is the position in buf.
	// If offset > 0, buf[offset:] has bytes which are read but haven't used.
	offset int
	// nextMsgLen is the length of the next log message.
	// If nextMsgLen = 0, a new value must be read from rdr.
	nextMsgLen int
}

func (d *decoder) readRecord(size int) error {
	var err error
	for i := 0; i < maxDecodeRetry; i++ {
		var n int
		n, err = io.ReadFull(d.rdr, d.buf[d.offset:size])
		d.offset += n
		if err != nil {
			if err != io.ErrUnexpectedEOF {
				return err
			}
			continue
		}
		break
	}
	if err != nil {
		return err
	}
	d.offset = 0
	return nil
}

func (d *decoder) Decode() (*logger.Message, error) {
	if d.proto == nil {
		d.proto = &logdriver.LogEntry{}
	} else {
		resetProto(d.proto)
	}
	if d.buf == nil {
		d.buf = make([]byte, initialBufSize)
	}

	if d.nextMsgLen == 0 {
		msgLen, err := d.decodeSizeHeader()
		if err != nil {
			return nil, err
		}

		if msgLen > maxMsgLen {
			return nil, fmt.Errorf("log message is too large (%d > %d)", msgLen, maxMsgLen)
		}

		if len(d.buf) < msgLen+encodeBinaryLen {
			d.buf = make([]byte, msgLen+encodeBinaryLen)
		} else if msgLen <= initialBufSize {
			d.buf = d.buf[:initialBufSize]
		} else {
			d.buf = d.buf[:msgLen+encodeBinaryLen]
		}

		d.nextMsgLen = msgLen
	}
	return d.decodeLogEntry()
}

func (d *decoder) Reset(rdr io.Reader) {
	if d.rdr == rdr {
		return
	}

	d.rdr = rdr
	if d.proto != nil {
		resetProto(d.proto)
	}
	if d.buf != nil {
		d.buf = d.buf[:initialBufSize]
	}
	d.offset = 0
	d.nextMsgLen = 0
}

func (d *decoder) Close() {
	d.buf = d.buf[:0]
	d.buf = nil
	if d.proto != nil {
		resetProto(d.proto)
	}
	d.rdr = nil
}

func decodeFunc(rdr io.Reader) loggerutils.Decoder {
	return &decoder{rdr: rdr}
}

func (d *decoder) decodeSizeHeader() (int, error) {
	err := d.readRecord(encodeBinaryLen)
	if err != nil {
		return 0, errors.Wrap(err, "could not read a size header")
	}

	msgLen := int(binary.BigEndian.Uint32(d.buf[:encodeBinaryLen]))
	return msgLen, nil
}

func (d *decoder) decodeLogEntry() (*logger.Message, error) {
	msgLen := d.nextMsgLen
	err := d.readRecord(msgLen + encodeBinaryLen)
	if err != nil {
		return nil, errors.Wrapf(err, "could not read a log entry (size=%d+%d)", msgLen, encodeBinaryLen)
	}
	d.nextMsgLen = 0

	if err := d.proto.Unmarshal(d.buf[:msgLen]); err != nil {
		return nil, errors.Wrapf(err, "error unmarshalling log entry (size=%d)", msgLen)
	}

	msg := protoToMessage(d.proto)
	if msg.PLogMetaData == nil || msg.PLogMetaData.Last {
		msg.Line = append(msg.Line, '\n')
	}

	return msg, nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2012, Suryandaru Triandana <syndtr@gmail.com>
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// +build darwin freebsd linux netbsd openbsd

package storage

import (
	"os"
	"syscall"
)

type unixFileLock struct {
	f *os.File
}

func (fl *unixFileLock) release() error {
	if err := setFileLock(fl.f, false); err != nil {
		return err
	}
	return fl.f.Close()
}

func newFileLock(path string) (fl fileLock, err error) {
	f, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0644)
	if err != nil {
		return
	}
	err = setFileLock(f, true)
	if err != nil {
		f.Close()
		return
	}
	fl = &unixFileLock{f: f}
	return
}

func setFileLock(f *os.File, lock bool) error {
	how := syscall.LOCK_UN
	if lock {
		how = syscall.LOCK_EX
	}
	return syscall.Flock(int(f.Fd()), how|syscall.LOCK_NB)
}

func rename(oldpath, newpath string) error {
	return os.Rename(oldpath, newpath)
}

func syncDir(name string) error {
	f, err := os.Open(name)
	if err != nil {
		return err
	}
	defer f.Close()
	if err := f.Sync(); err != nil {
		return err
	}
	return nil
}
<commit_msg>storage: Add support for dragonfly #59<commit_after>// Copyright (c) 2012, Suryandaru Triandana <syndtr@gmail.com>
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// +build darwin dragonfly freebsd linux netbsd openbsd

package storage

import (
	"os"
	"syscall"
)

type unixFileLock struct {
	f *os.File
}

func (fl *unixFileLock) release() error {
	if err := setFileLock(fl.f, false); err != nil {
		return err
	}
	return fl.f.Close()
}

func newFileLock(path string) (fl fileLock, err error) {
	f, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0644)
	if err != nil {
		return
	}
	err = setFileLock(f, true)
	if err != nil {
		f.Close()
		return
	}
	fl = &unixFileLock{f: f}
	return
}

func setFileLock(f *os.File, lock bool) error {
	how := syscall.LOCK_UN
	if lock {
		how = syscall.LOCK_EX
	}
	return syscall.Flock(int(f.Fd()), how|syscall.LOCK_NB)
}

func rename(oldpath, newpath string) error {
	return os.Rename(oldpath, newpath)
}

func syncDir(name string) error {
	f, err := os.Open(name)
	if err != nil {
		return err
	}
	defer f.Close()
	if err := f.Sync(); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"

	"github.com/hyperledger/fabric/core/chaincode/shim"
)

type SimpleChaincode struct {
}

func main() {

	err := shim.Start(new(SimpleChaincode))
	if err != nil {
		fmt.Printf("Error starting Simple chaincode: %s", err)
	}
}

func (t *SimpleChaincode) Init(stub shim.ChaincodeStubInterface, function string, args []string) ([]byte, error) {
	if len(args) != 1 {
		return nil, errors.New("incorrect number of args")
	}

	err := stub.PutState("hello_world", []byte(args[0]))
	if err != nil {
		return nil, err
	}

	return nil, nil
}

func (t *SimpleChaincode) Invoke(stub shim.ChaincodeStubInterface, function string, args []string) ([]byte, error) {
	fmt.Println("invoke is running" + function)

	if function == "init" {
		return t.Init(stub, "init", args)
	} else if function == "write" {
		return t.write(stub, args)
	}
	fmt.Println("invoke did not find func:" + function)

	return nil, errors.New("received unknow function" + function)
}

func (t *SimpleChaincode) Query(stub shim.ChaincodeStubInterface, function string, args []string) ([]byte, error) {
	fmt.Println("query is running " + function) //log
	switch function {
	case "read":
		return t.read(stub, args)
	case "read2":
		return t.read2(stub, args)
	case "helloworld":
		return t.helloworld()
	default:
		fmt.Println("query did not find func: " + function)
	} //根据输入的string类型参数function, 来决定使用哪个子函数。

	return nil, errors.New("received unknow function: " + function) //如果没有的话，则返回失败
}

func (t *SimpleChaincode) read2(stub shim.ChaincodeStubInterface, args []string) ([]byte, error) {
	key := args[0]
	result := []byte("helloworld" + key)
	var err error
	return result, err
}

func (t *SimpleChaincode) read3(stub shim.ChaincodeStubInterface) ([]byte, error) {
	result := []byte(stub.GetTxID())
	return result, nil
}

func (t *SimpleChaincode) helloworld() ([]byte, error) {
	result := []byte("hello world")
	return result, nil
}

func (t *SimpleChaincode) write(stub shim.ChaincodeStubInterface, args []string) ([]byte, error) {
	var key, value string
	var err error
	fmt.Println("running write()")

	if len(args) != 2 {
		return nil, errors.New("Incorrect number of arguments. Expecting 2")
	}

	key = args[0]
	value = args[1]

	err = stub.PutState(key, []byte(value))

	if err != nil {
		return nil, err
	}

	return nil, nil
}

func (t *SimpleChaincode) read(stub shim.ChaincodeStubInterface, args []string) ([]byte, error) {
	var key, jsonResp string
	var err error

	if len(args) != 1 {
		return nil, errors.New("incorrect number of arguments. expecting 1")
	}

	key = args[0]
	valAsbytes, err := stub.GetState(key)

	if err != nil {
		jsonResp = "{\"Error\":\"Failed to get state for " + key + "\"}"
		return nil, errors.New(jsonResp)
	}

	return valAsbytes, nil
}
<commit_msg>fix bugs<commit_after>package main

import (
	"errors"
	"fmt"

	"github.com/hyperledger/fabric/core/chaincode/shim"
)

type SimpleChaincode struct {
}

func main() {

	err := shim.Start(new(SimpleChaincode))
	if err != nil {
		fmt.Printf("Error starting Simple chaincode: %s", err)
	}
}

func (t *SimpleChaincode) Init(stub shim.ChaincodeStubInterface, function string, args []string) ([]byte, error) {
	if len(args) != 1 {
		return nil, errors.New("incorrect number of args")
	}

	err := stub.PutState("hello_world", []byte(args[0]))
	if err != nil {
		return nil, err
	}

	return nil, nil
}

func (t *SimpleChaincode) Invoke(stub shim.ChaincodeStubInterface, function string, args []string) ([]byte, error) {
	fmt.Println("invoke is running" + function)

	if function == "init" {
		return t.Init(stub, "init", args)
	} else if function == "write" {
		return t.write(stub, args)
	}
	fmt.Println("invoke did not find func:" + function)

	return nil, errors.New("received unknow function" + function)
}

func (t *SimpleChaincode) Query(stub shim.ChaincodeStubInterface, function string, args []string) ([]byte, error) {
	fmt.Println("query is running " + function) //log
	switch function {
	case "read":
		return t.read(stub, args)
	case "read2":
		return t.read2(stub, args)
	case "read3":
		return t.read3(stub)
	case "helloworld":
		return t.helloworld()
	default:
		fmt.Println("query did not find func: " + function)
	} //根据输入的string类型参数function, 来决定使用哪个子函数。

	return nil, errors.New("received unknow function: " + function) //如果没有的话，则返回失败
}

func (t *SimpleChaincode) read2(stub shim.ChaincodeStubInterface, args []string) ([]byte, error) {
	key := args[0]
	result := []byte("helloworld" + key)
	var err error
	return result, err
}

func (t *SimpleChaincode) read3(stub shim.ChaincodeStubInterface) ([]byte, error) {
	result := []byte(stub.GetTxID())
	return result, nil
}

func (t *SimpleChaincode) helloworld() ([]byte, error) {
	result := []byte("hello world")
	return result, nil
}

func (t *SimpleChaincode) read(stub shim.ChaincodeStubInterface, args []string) ([]byte, error) {
	var key, jsonResp string
	var err error
	if len(args) != 1 {
		return nil, errors.New("incorrect number of arguments. expecting 1")
	}
	key = args[0]
	valAsbytes, err := stub.GetState(key)

	if err != nil {
		jsonResp = "{\"Error\":\"Failed to get state for " + key + "\"}"
		return nil, errors.New(jsonResp)
	}

	return valAsbytes, nil
}

func (t *SimpleChaincode) write(stub shim.ChaincodeStubInterface, args []string) ([]byte, error) {
	var key, value string
	var err error
	fmt.Println("running write()")

	if len(args) != 2 {
		return nil, errors.New("Incorrect number of arguments. Expecting 2")
	}

	key = args[0]
	value = args[1]

	err = stub.PutState(key, []byte(value))

	if err != nil {
		return nil, err
	}

	return nil, nil
}
<|endoftext|>


<commit_before>package v1alpha1

import (
	crdutils "github.com/appscode/kutil/apiextensions/v1beta1"
	apiextensions "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1beta1"
)

func (d DormantDatabase) OffshootName() string {
	return d.Name
}

func (d DormantDatabase) ResourceShortCode() string {
	return ResourceCodeDormantDatabase
}

func (d DormantDatabase) ResourceKind() string {
	return ResourceKindDormantDatabase
}

func (d DormantDatabase) ResourceSingular() string {
	return ResourceSingularDormantDatabase
}

func (d DormantDatabase) ResourcePlural() string {
	return ResourcePluralDormantDatabase
}

func (d DormantDatabase) CustomResourceDefinition() *apiextensions.CustomResourceDefinition {
	return crdutils.NewCustomResourceDefinition(crdutils.Config{
		Group:         SchemeGroupVersion.Group,
		Plural:        ResourcePluralDormantDatabase,
		Singular:      ResourceSingularDormantDatabase,
		Kind:          ResourceKindDormantDatabase,
		ShortNames:    []string{ResourceCodeDormantDatabase},
		Categories:    []string{"datastore", "kubedb", "appscode", "all"},
		ResourceScope: string(apiextensions.NamespaceScoped),
		Versions: []apiextensions.CustomResourceDefinitionVersion{
			{
				Name:    SchemeGroupVersion.Version,
				Served:  true,
				Storage: true,
			},
		},
		Labels: crdutils.Labels{
			LabelsMap: map[string]string{"app": "kubedb"},
		},
		SpecDefinitionName:      "github.com/kubedb/apimachinery/apis/kubedb/v1alpha1.DormantDatabase",
		EnableValidation:        false,
		GetOpenAPIDefinitions:   GetOpenAPIDefinitions,
		EnableStatusSubresource: EnableStatusSubresource,
		AdditionalPrinterColumns: []apiextensions.CustomResourceColumnDefinition{
			{
				Name:     "Status",
				Type:     "string",
				JSONPath: ".status.phase",
			},
			{
				Name:     "Age",
				Type:     "date",
				JSONPath: ".metadata.creationTimestamp",
			},
		},
	}, setNameSchema)
}

func (d *DormantDatabase) Migrate() {
	if d == nil {
		return
	}
	d.Spec.Origin.Spec.Elasticsearch.Migrate()
	d.Spec.Origin.Spec.Postgres.Migrate()
	d.Spec.Origin.Spec.MySQL.Migrate()
	d.Spec.Origin.Spec.MongoDB.Migrate()
	d.Spec.Origin.Spec.Redis.Migrate()
	d.Spec.Origin.Spec.Memcached.Migrate()
	d.Spec.Origin.Spec.Etcd.Migrate()
}
<commit_msg>Add OffshootSelectors to DormantDatabase (#294)<commit_after>package v1alpha1

import (
	crdutils "github.com/appscode/kutil/apiextensions/v1beta1"
	meta_util "github.com/appscode/kutil/meta"
	apiextensions "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1beta1"
)

func (d DormantDatabase) OffshootSelectors() map[string]string {
	selector := map[string]string{
		LabelDatabaseName: d.Name,
	}
	switch {
	case d.Spec.Origin.Spec.Etcd != nil:
		selector[LabelDatabaseKind] = ResourceKindEtcd
	case d.Spec.Origin.Spec.Elasticsearch != nil:
		selector[LabelDatabaseKind] = ResourceKindElasticsearch
	case d.Spec.Origin.Spec.Memcached != nil:
		selector[LabelDatabaseKind] = ResourceKindMemcached
	case d.Spec.Origin.Spec.MongoDB != nil:
		selector[LabelDatabaseKind] = ResourceKindMongoDB
	case d.Spec.Origin.Spec.MySQL != nil:
		selector[LabelDatabaseKind] = ResourceKindMySQL
	case d.Spec.Origin.Spec.Postgres != nil:
		selector[LabelDatabaseKind] = ResourceKindPostgres
	case d.Spec.Origin.Spec.Redis != nil:
		selector[LabelDatabaseKind] = ResourceKindRedis
	}
	return selector
}

func (d DormantDatabase) OffshootLabels() map[string]string {
	return meta_util.FilterKeys(GenericKey, d.OffshootSelectors(), d.Spec.Origin.Labels)
}

func (d DormantDatabase) OffshootName() string {
	return d.Name
}

func (d DormantDatabase) ResourceShortCode() string {
	return ResourceCodeDormantDatabase
}

func (d DormantDatabase) ResourceKind() string {
	return ResourceKindDormantDatabase
}

func (d DormantDatabase) ResourceSingular() string {
	return ResourceSingularDormantDatabase
}

func (d DormantDatabase) ResourcePlural() string {
	return ResourcePluralDormantDatabase
}

func (d DormantDatabase) CustomResourceDefinition() *apiextensions.CustomResourceDefinition {
	return crdutils.NewCustomResourceDefinition(crdutils.Config{
		Group:         SchemeGroupVersion.Group,
		Plural:        ResourcePluralDormantDatabase,
		Singular:      ResourceSingularDormantDatabase,
		Kind:          ResourceKindDormantDatabase,
		ShortNames:    []string{ResourceCodeDormantDatabase},
		Categories:    []string{"datastore", "kubedb", "appscode", "all"},
		ResourceScope: string(apiextensions.NamespaceScoped),
		Versions: []apiextensions.CustomResourceDefinitionVersion{
			{
				Name:    SchemeGroupVersion.Version,
				Served:  true,
				Storage: true,
			},
		},
		Labels: crdutils.Labels{
			LabelsMap: map[string]string{"app": "kubedb"},
		},
		SpecDefinitionName:      "github.com/kubedb/apimachinery/apis/kubedb/v1alpha1.DormantDatabase",
		EnableValidation:        false,
		GetOpenAPIDefinitions:   GetOpenAPIDefinitions,
		EnableStatusSubresource: EnableStatusSubresource,
		AdditionalPrinterColumns: []apiextensions.CustomResourceColumnDefinition{
			{
				Name:     "Status",
				Type:     "string",
				JSONPath: ".status.phase",
			},
			{
				Name:     "Age",
				Type:     "date",
				JSONPath: ".metadata.creationTimestamp",
			},
		},
	}, setNameSchema)
}

func (d *DormantDatabase) Migrate() {
	if d == nil {
		return
	}
	d.Spec.Origin.Spec.Elasticsearch.Migrate()
	d.Spec.Origin.Spec.Postgres.Migrate()
	d.Spec.Origin.Spec.MySQL.Migrate()
	d.Spec.Origin.Spec.MongoDB.Migrate()
	d.Spec.Origin.Spec.Redis.Migrate()
	d.Spec.Origin.Spec.Memcached.Migrate()
	d.Spec.Origin.Spec.Etcd.Migrate()
}
<|endoftext|>


<commit_before>package reservation

import (
	"fmt"
	"os"
	"time"

	"github.com/garyburd/redigo/redis"
)

// Reservation is a type that represents a lock on a resource. At most one reservation
// can exist for an individual resource at any time.
type Reservation struct {
	stopped     bool
	key, source string
	getConn     func() redis.Conn
	ttl         time.Duration
}

// Manager is responsible for creating and extending reservations. When a Reservation
// is created using Manager, the manager will automatically extend that Reservation
// every `Manager.Heartbeat` time units by setting the Reservation to expire
// after `Manager.TTL` time elapses.
type Manager struct {
	owner          string
	pool           *redis.Pool
	Heartbeat, TTL time.Duration
}

// NewManager returns a new Manager, or an error if a connection to the supplied
// Redis server cannot be made.
func NewManager(redisURL, owner string) (*Manager, error) {
	// Open redis pool
	redisPool := redis.NewPool(func() (redis.Conn, error) {
		return redis.DialTimeout("tcp", redisURL, 15*time.Second, 10*time.Second, 10*time.Second)
	}, 5)

	// Get a conn and ping so we fail immediately if the URL is wrong
	conn := redisPool.Get()
	defer conn.Close()
	if _, err := conn.Do("PING"); err != nil {
		return nil, fmt.Errorf("Error connecting to redis: %s", err)
	}

	return &Manager{
		Heartbeat: 15 * time.Minute,
		TTL:       4 * time.Hour,
		owner:     owner,
		pool:      redisPool,
	}, nil
}

// Lock creates a Reservation for `resource`, or returns an error if there already exists a
// Reservation for that resource.
func (manager *Manager) Lock(resource string) (*Reservation, error) {
	// Get hostname
	hostname, err := os.Hostname()
	if err != nil {
		return nil, err
	}

	key := fmt.Sprintf("reservation-%s", resource)
	val := fmt.Sprintf("%s-%s-%d", hostname, manager.owner, os.Getpid())

	// Get connection
	conn := manager.pool.Get()
	defer conn.Close()

	// Try to set the reservation
	success, err := conn.Do(
		"SET", key, val,
		"EX", manager.TTL.Seconds(),
		"NX")
	if err != nil {
		return nil, fmt.Errorf("Error with SET command: %s", err.Error())
	}
	if success == nil {
		return nil, fmt.Errorf("Reservation already exists for resource %s", resource)
	}

	// Make new reservation
	res := &Reservation{
		key:     key,
		source:  val,
		getConn: manager.pool.Get,
		ttl:     manager.TTL,
	}

	// Set up heartbeat in background
	go func() {
		for _ = range time.Tick(manager.Heartbeat) {
			if res.stopped {
				break
			}
			// Panic if err; no way to handle the error gracefully when this runs in the background
			if _, err := res.heartbeat(); err != nil {
				panic(err)
			}
		}
	}()

	return res, nil
}

// Release ends a lock on a resource. Release returns `nil` if release was successful or
// an `error` if not. In the event of an error, the reservation will be removed from Redis after
// `Reservation.ttl` expires.
func (res *Reservation) Release() error {
	conn := res.getConn()
	defer conn.Close()

	_, err := redis.Int(conn.Do("DEL", res.key))
	// Always release lock so reservation will expire after TTL if delete fails
	res.stopped = true

	if err != nil {
		return fmt.Errorf("Error deleting reservation key for %s: %s", res.key, err.Error())
	}
	return nil
}

func (res *Reservation) heartbeat() (int, error) {
	// Get connection
	conn := res.getConn()
	defer conn.Close()

	// Check that the reservation still exists and error if we don't have it
	source, err := redis.String(conn.Do("GET", res.key))
	if err != nil {
		return -1, fmt.Errorf("Could not fetch owner of reservation %s: ERR %s", res.key, err.Error())
	}
	if source != res.source {
		return -1, fmt.Errorf("Reservation for %s has unknown owner %s", res.key, source)
	}

	// Extend reservation
	success, err := redis.Int(conn.Do("EXPIRE", res.key, res.ttl.Seconds()))
	if err != nil || success != 1 {
		return -1, fmt.Errorf("Could not extend reservation %s: ERR %s", res.key, err.Error())
	}
	return success, nil
}
<commit_msg>handle redis errors safely<commit_after>package reservation

import (
	"fmt"
	"os"
	"time"

	"github.com/garyburd/redigo/redis"
)

// Reservation is a type that represents a lock on a resource. At most one reservation
// can exist for an individual resource at any time.
type Reservation struct {
	stopped     bool
	key, source string
	getConn     func() redis.Conn
	ttl         time.Duration
}

// Manager is responsible for creating and extending reservations. When a Reservation
// is created using Manager, the manager will automatically extend that Reservation
// every `Manager.Heartbeat` time units by setting the Reservation to expire
// after `Manager.TTL` time elapses.
type Manager struct {
	owner          string
	pool           *redis.Pool
	Heartbeat, TTL time.Duration
}

// NewManager returns a new Manager, or an error if a connection to the supplied
// Redis server cannot be made.
func NewManager(redisURL, owner string) (*Manager, error) {
	// Open redis pool
	redisPool := redis.NewPool(func() (redis.Conn, error) {
		return redis.DialTimeout("tcp", redisURL, 15*time.Second, 10*time.Second, 10*time.Second)
	}, 5)

	// Get a conn and ping so we fail immediately if the URL is wrong
	conn := redisPool.Get()
	defer conn.Close()
	if _, err := conn.Do("PING"); err != nil {
		return nil, fmt.Errorf("Error connecting to redis: %s", err)
	}

	return &Manager{
		Heartbeat: 15 * time.Minute,
		TTL:       4 * time.Hour,
		owner:     owner,
		pool:      redisPool,
	}, nil
}

// Lock creates a Reservation for `resource`, or returns an error if there already exists a
// Reservation for that resource.
func (manager *Manager) Lock(resource string) (*Reservation, error) {
	// Get hostname
	hostname, err := os.Hostname()
	if err != nil {
		return nil, err
	}

	key := fmt.Sprintf("reservation-%s", resource)
	val := fmt.Sprintf("%s-%s-%d", hostname, manager.owner, os.Getpid())

	// Get connection
	conn := manager.pool.Get()
	defer conn.Close()

	// Try to set the reservation
	success, err := conn.Do(
		"SET", key, val,
		"EX", manager.TTL.Seconds(),
		"NX")
	if err != nil {
		return nil, fmt.Errorf("Error with SET command: %s", err.Error())
	}
	if success == nil {
		return nil, fmt.Errorf("Reservation already exists for resource %s", resource)
	}

	// Make new reservation
	res := &Reservation{
		key:     key,
		source:  val,
		getConn: manager.pool.Get,
		ttl:     manager.TTL,
	}

	// Set up heartbeat in background
	go func() {
		for _ = range time.Tick(manager.Heartbeat) {
			if res.stopped {
				break
			}
			// Panic if err; no way to handle the error gracefully when this runs in the background
			success, err := res.heartbeat()
			if err != nil {
				panic(err)
			}
			if success != 1 {
				panic(fmt.Errorf("Got code %d when attempting to extend reservation", success))
			}
		}
	}()

	return res, nil
}

// Release ends a lock on a resource. Release returns `nil` if release was successful or
// an `error` if not. In the event of an error, the reservation will be removed from Redis after
// `Reservation.ttl` expires.
func (res *Reservation) Release() error {
	conn := res.getConn()
	defer conn.Close()

	_, err := redis.Int(conn.Do("DEL", res.key))
	// Always release lock so reservation will expire after TTL if delete fails
	res.stopped = true

	if err != nil {
		return fmt.Errorf("Error deleting reservation key for %s: %s", res.key, err.Error())
	}
	return nil
}

func (res *Reservation) heartbeat() (int, error) {
	// Get connection
	conn := res.getConn()
	defer conn.Close()

	// Check that the reservation still exists and error if we don't have it
	source, err := redis.String(conn.Do("GET", res.key))
	if err != nil {
		return -1, fmt.Errorf("Could not fetch owner of reservation %s: ERR %s", res.key, err.Error())
	}
	if source != res.source {
		return -1, fmt.Errorf("Reservation for %s has unknown owner %s", res.key, source)
	}

	// Extend reservation
	success, err := redis.Int(conn.Do("EXPIRE", res.key, res.ttl.Seconds()))
	if err != nil {
		return -1, fmt.Errorf("Could not extend reservation %s: ERR %s", res.key, err.Error())
	}
	return success, nil
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/cloudtrail"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsCloudTrail() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsCloudTrailCreate,
		Read:   resourceAwsCloudTrailRead,
		Update: resourceAwsCloudTrailUpdate,
		Delete: resourceAwsCloudTrailDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"s3_bucket_name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"s3_key_prefix": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"cloud_watch_logs_role_arn": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"cloud_watch_logs_group_arn": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"include_global_service_events": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  true,
			},
			"sns_topic_name": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
		},
	}
}

func resourceAwsCloudTrailCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).cloudtrailconn

	input := cloudtrail.CreateTrailInput{
		Name:         aws.String(d.Get("name").(string)),
		S3BucketName: aws.String(d.Get("s3_bucket_name").(string)),
	}

	if v, ok := d.GetOk("cloud_watch_logs_group_arn"); ok {
		input.CloudWatchLogsLogGroupArn = aws.String(v.(string))
	}
	if v, ok := d.GetOk("cloud_watch_logs_role_arn"); ok {
		input.CloudWatchLogsRoleArn = aws.String(v.(string))
	}
	if v, ok := d.GetOk("include_global_service_events"); ok {
		input.IncludeGlobalServiceEvents = aws.Bool(v.(bool))
	}
	if v, ok := d.GetOk("s3_key_prefix"); ok {
		input.S3KeyPrefix = aws.String(v.(string))
	}
	if v, ok := d.GetOk("sns_topic_name"); ok {
		input.SnsTopicName = aws.String(v.(string))
	}

	t, err := conn.CreateTrail(&input)
	if err != nil {
		return err
	}

	log.Printf("[DEBUG] CloudTrail created: %s", t)

	d.SetId(*t.Name)

	return resourceAwsCloudTrailRead(d, meta)
}

func resourceAwsCloudTrailRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).cloudtrailconn

	name := d.Get("name").(string)
	input := cloudtrail.DescribeTrailsInput{
		TrailNameList: []*string{
			aws.String(name),
		},
	}
	resp, err := conn.DescribeTrails(&input)
	if err != nil {
		return err
	}
	if len(resp.TrailList) == 0 {
		return fmt.Errorf("No CloudTrail found, using name %q", name)
	}

	trail := resp.TrailList[0]
	log.Printf("[DEBUG] CloudTrail received: %s", trail)

	d.Set("name", trail.Name)
	d.Set("s3_bucket_name", trail.S3BucketName)
	d.Set("s3_key_prefix", trail.S3KeyPrefix)
	d.Set("cloud_watch_logs_role_arn", trail.CloudWatchLogsRoleArn)
	d.Set("cloud_watch_logs_group_arn", trail.CloudWatchLogsLogGroupArn)
	d.Set("include_global_service_events", trail.IncludeGlobalServiceEvents)
	d.Set("sns_topic_name", trail.SnsTopicName)

	return nil
}

func resourceAwsCloudTrailUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).cloudtrailconn

	input := cloudtrail.UpdateTrailInput{
		Name: aws.String(d.Get("name").(string)),
	}

	if d.HasChange("s3_bucket_name") {
		input.S3BucketName = aws.String(d.Get("s3_bucket_name").(string))
	}
	if d.HasChange("s3_key_prefix") {
		input.S3KeyPrefix = aws.String(d.Get("s3_key_prefix").(string))
	}
	if d.HasChange("cloud_watch_logs_role_arn") {
		input.CloudWatchLogsRoleArn = aws.String(d.Get("cloud_watch_logs_role_arn").(string))
	}
	if d.HasChange("cloud_watch_logs_group_arn") {
		input.CloudWatchLogsLogGroupArn = aws.String(d.Get("cloud_watch_logs_group_arn").(string))
	}
	if d.HasChange("include_global_service_events") {
		input.IncludeGlobalServiceEvents = aws.Bool(d.Get("include_global_service_events").(bool))
	}
	if d.HasChange("sns_topic_name") {
		input.SnsTopicName = aws.String(d.Get("sns_topic_name").(string))
	}

	log.Printf("[DEBUG] Updating CloudTrail: %s", input)
	t, err := conn.UpdateTrail(&input)
	if err != nil {
		return err
	}
	log.Printf("[DEBUG] CloudTrail updated: %s", t)

	return resourceAwsCloudTrailRead(d, meta)
}

func resourceAwsCloudTrailDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).cloudtrailconn
	name := d.Get("name").(string)

	log.Printf("[DEBUG] Deleting CloudTrail: %q", name)
	_, err := conn.DeleteTrail(&cloudtrail.DeleteTrailInput{
		Name: aws.String(name),
	})

	return err
}
<commit_msg>Add enable_logging to AWS CloudTrail<commit_after>package aws

import (
	"fmt"
	"log"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/cloudtrail"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsCloudTrail() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsCloudTrailCreate,
		Read:   resourceAwsCloudTrailRead,
		Update: resourceAwsCloudTrailUpdate,
		Delete: resourceAwsCloudTrailDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"enable_logging": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  false,
			},
			"s3_bucket_name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"s3_key_prefix": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"cloud_watch_logs_role_arn": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"cloud_watch_logs_group_arn": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"include_global_service_events": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  true,
			},
			"sns_topic_name": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
		},
	}
}

func resourceAwsCloudTrailCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).cloudtrailconn

	input := cloudtrail.CreateTrailInput{
		Name:         aws.String(d.Get("name").(string)),
		S3BucketName: aws.String(d.Get("s3_bucket_name").(string)),
	}

	if v, ok := d.GetOk("cloud_watch_logs_group_arn"); ok {
		input.CloudWatchLogsLogGroupArn = aws.String(v.(string))
	}
	if v, ok := d.GetOk("cloud_watch_logs_role_arn"); ok {
		input.CloudWatchLogsRoleArn = aws.String(v.(string))
	}
	if v, ok := d.GetOk("include_global_service_events"); ok {
		input.IncludeGlobalServiceEvents = aws.Bool(v.(bool))
	}
	if v, ok := d.GetOk("s3_key_prefix"); ok {
		input.S3KeyPrefix = aws.String(v.(string))
	}
	if v, ok := d.GetOk("sns_topic_name"); ok {
		input.SnsTopicName = aws.String(v.(string))
	}

	t, err := conn.CreateTrail(&input)
	if err != nil {
		return err
	}

	log.Printf("[DEBUG] CloudTrail created: %s", t)

	d.SetId(*t.Name)

	// AWS CloudTrail sets newly-created trails to false.
	if v, ok := d.GetOk("enable_logging"); ok && v.(bool) {
		cloudTrailSetLogging(conn, v.(bool), d.Id())
	}

	return resourceAwsCloudTrailRead(d, meta)
}

func resourceAwsCloudTrailRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).cloudtrailconn

	name := d.Get("name").(string)
	input := cloudtrail.DescribeTrailsInput{
		TrailNameList: []*string{
			aws.String(name),
		},
	}
	resp, err := conn.DescribeTrails(&input)
	if err != nil {
		return err
	}
	if len(resp.TrailList) == 0 {
		return fmt.Errorf("No CloudTrail found, using name %q", name)
	}

	trail := resp.TrailList[0]
	log.Printf("[DEBUG] CloudTrail received: %s", trail)

	d.Set("name", trail.Name)
	d.Set("s3_bucket_name", trail.S3BucketName)
	d.Set("s3_key_prefix", trail.S3KeyPrefix)
	d.Set("cloud_watch_logs_role_arn", trail.CloudWatchLogsRoleArn)
	d.Set("cloud_watch_logs_group_arn", trail.CloudWatchLogsLogGroupArn)
	d.Set("include_global_service_events", trail.IncludeGlobalServiceEvents)
	d.Set("sns_topic_name", trail.SnsTopicName)

	logstatus, err := cloudTrailGetLoggingStatus(conn, *trail.Name)
	if err != nil {
		return err
	}
	d.Set("enable_logging", logstatus)

	return nil
}

func resourceAwsCloudTrailUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).cloudtrailconn

	input := cloudtrail.UpdateTrailInput{
		Name: aws.String(d.Get("name").(string)),
	}

	if d.HasChange("s3_bucket_name") {
		input.S3BucketName = aws.String(d.Get("s3_bucket_name").(string))
	}
	if d.HasChange("s3_key_prefix") {
		input.S3KeyPrefix = aws.String(d.Get("s3_key_prefix").(string))
	}
	if d.HasChange("cloud_watch_logs_role_arn") {
		input.CloudWatchLogsRoleArn = aws.String(d.Get("cloud_watch_logs_role_arn").(string))
	}
	if d.HasChange("cloud_watch_logs_group_arn") {
		input.CloudWatchLogsLogGroupArn = aws.String(d.Get("cloud_watch_logs_group_arn").(string))
	}
	if d.HasChange("include_global_service_events") {
		input.IncludeGlobalServiceEvents = aws.Bool(d.Get("include_global_service_events").(bool))
	}
	if d.HasChange("sns_topic_name") {
		input.SnsTopicName = aws.String(d.Get("sns_topic_name").(string))
	}

	log.Printf("[DEBUG] Updating CloudTrail: %s", input)
	t, err := conn.UpdateTrail(&input)
	if err != nil {
		return err
	}

	if d.HasChange("enable_logging") {
		log.Printf("[DEBUG] Updating logging on CloudTrail: %s", input)
		err := cloudTrailSetLogging(conn, d.Get("enable_logging").(bool), *input.Name)
		if err != nil {
			return err
		}
	}

	log.Printf("[DEBUG] CloudTrail updated: %s", t)

	return resourceAwsCloudTrailRead(d, meta)
}

func resourceAwsCloudTrailDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).cloudtrailconn
	name := d.Get("name").(string)

	log.Printf("[DEBUG] Deleting CloudTrail: %q", name)
	_, err := conn.DeleteTrail(&cloudtrail.DeleteTrailInput{
		Name: aws.String(name),
	})

	return err
}

func cloudTrailGetLoggingStatus(conn *cloudtrail.CloudTrail, id string) (bool, error) {
	GetTrailStatusOpts := &cloudtrail.GetTrailStatusInput{
		Name: aws.String(id),
	}
	resp, err := conn.GetTrailStatus(GetTrailStatusOpts)

	return *resp.IsLogging, err
}

func cloudTrailSetLogging(conn *cloudtrail.CloudTrail, enabled bool, id string) error {
	if enabled {
		log.Printf(
			"[DEBUG] Starting logging on CloudTrail (%s)",
			id)
		StartLoggingOpts := &cloudtrail.StartLoggingInput{
			Name: aws.String(id),
		}
		if _, err := conn.StartLogging(StartLoggingOpts); err != nil {
			return fmt.Errorf(
				"Error starting logging on CloudTrail (%s): %s",
				id, err)
		}
	} else {
		log.Printf(
			"[DEBUG] Stopping logging on CloudTrail (%s)",
			id)
		StopLoggingOpts := &cloudtrail.StopLoggingInput{
			Name: aws.String(id),
		}
		if _, err := conn.StopLogging(StopLoggingOpts); err != nil {
			return fmt.Errorf(
				"Error stopping logging on CloudTrail (%s): %s",
				id, err)
		}
	}

	return nil
}
<|endoftext|>