<commit_before>// Copyright (c) 2013, Bartosz Kratochwil
// All rights reserved.

// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:

//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.

//     * Redistributions in binary form must reproduce the above
//       copyright notice, this list of conditions and the following
//       disclaimer in the documentation and/or other materials provided
//       with the distribution.

//     * Neither the name of Aaron France nor the names of its
//       contributors may be used to endorse or promote products derived
//       from this software without specific prior written permission.

// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
package hpcloud

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
)

/*
 ListDBInstances will list all the available database instances

This function implements the interface as described in:
http://api-docs.hpcloud.com/hpcloud-rdb-mysql/1.0/content/list-database-instances.html
*/
func (a Access) ListDBInstances() (*DBInstances, error) {
	url := fmt.Sprintf("%s%s/instances", RDB_URL, a.TenantID)
	body, err := a.baseRequest(url, "GET", nil)
	if err != nil {
		return nil, err
	}
	dbs := &DBInstances{}
	err = json.Unmarshal(body, dbs)
	return dbs, err
}

/*
 This function takes instance ID and deletes database instance with this ID.

This function implements the interface as described in:
http://api-docs.hpcloud.com/hpcloud-rdb-mysql/1.0/content/delete-instance.html
*/
func (a Access) DeleteDBInstance(instanceID string) error {
	url := fmt.Sprintf("%s%s/instances/%s", RDB_URL, a.TenantID,
		instanceID)
	_, err := a.baseRequest(url, "DELETE", nil)
	return err
}

/*
 This function takes instance ID and restarts DB instance with this ID.

This function implements the interface as described in:
http://api-docs.hpcloud.com/hpcloud-rdb-mysql/1.0/content/restart-instance.html
*/
func (a Access) RestartDBInstance(instanceID string) error {
	b := `{"restart":{}}`
	url := fmt.Sprintf("%s%s/instances/%s/action", RDB_URL,
		a.TenantID, instanceID)
	_, err := a.baseRequest(url, "POST", strings.NewReader(b))
	return err
}

/*
 ListAllFlavors lists all available database flavors.
 This function implements interface as described in:-
http://api-docs.hpcloud.com/hpcloud-rdb-mysql/1.0/content/list-flavors.html
*/
func (a Access) ListAllFlavors() (*DBFlavors, error) {
	url := fmt.Sprintf("%s%s/flavors", RDB_URL, a.TenantID)
	body, err := a.baseRequest(url, "GET", nil)
	if err != nil {
		return nil, err
	}

	flv := &DBFlavors{}
	err = json.Unmarshal(body, flv)
	return flv, err
}

/*
 This function returns flavor specs for given flavor.

 This function implements the interface as described in:
 http://api-docs.hpcloud.com/hpcloud-rdb-mysql/1.0/content/get-flavor.html
*/
func (a Access) GetDBFlavor(ID string) (*DBFlavor, error) {
	url := fmt.Sprintf("%s/flavors/%s", RDB_URL, ID)
	body, err := a.baseRequest(url, "GET", nil)
	if err != nil {
		return nil, err
	}

	flv := &DBFlavor{}
	err = json.Unmarshal(body, flv)
	if err != nil {
		return nil, err
	}
	return flv, nil
}

/*
 CreateDBInstance creates new database instance in the HPCloud using
settings found in the DatabaseReq instance passed to this function

 This function implements the interface as described in:
 http://api-docs.hpcloud.com/hpcloud-rdb-mysql/1.0/content/create-instance.html
*/
func (a Access) CreateDBInstance(db DatabaseReq) (*NewDBInstance, error) {
	b, err := db.MarshalJSON()
	if err != nil {
		return nil, err
	}

	url := fmt.Sprintf("%s%s/instances", RDB_URL, a.TenantID)

	body, err := a.baseRequest(url, "POST",
		strings.NewReader(string(b)))
	if err != nil {
		return nil, err
	}

	type respDB struct {
		Instance NewDBInstance `json:"instance"`
	}

	sr := &respDB{}
	err = json.Unmarshal(body, sr)
	if err != nil {
		return nil, err
	}
	return &sr.Instance, nil
}

/*
 This function retrieves details of the instance with provided ID.

This function implements the interface as described in:
http://api-docs.hpcloud.com/hpcloud-rdb-mysql/1.0/content/get-instance.html
*/
func (a Access) GetDBInstance(id string) (*InstDetails, error) {
	url := fmt.Sprintf("%s%s/instances/%s", RDB_URL, a.TenantID, id)
	body, err := a.baseRequest(url, "GET", nil)
	if err != nil {
		return nil, err
	}
	det := &InstDetails{}
	err = json.Unmarshal(body, det)
	return det, err
}

/*
 This function takes instance ID and resets password for this instance. It
 returns a new instance password.

 This function implements the interface as decribed in:
 http://api-docs.hpcloud.com/hpcloud-rdb-mysql/1.0/content/reset-instance-password.html
*/
func (a Access) ResetDBPassword(id string) (*DBCredentials, error) {
	b := `{"reset-password":{}}`
	url := fmt.Sprintf("%s%s/instances/%s/action", RDB_URL,
		a.TenantID, id)
	body, err := a.baseRequest(url, "POST", strings.NewReader(b))

	sr := &DBCredentials{}
	err = json.Unmarshal(body, sr)
	if err != nil {
		return nil, err
	}
	return sr, nil
}

/*
 This function lists all the security groups available for tenant.

 This function implements the interface as described in:
 http://api-docs.hpcloud.com/hpcloud-rdb-mysql/1.0/content/list-security-groups.html
*/
func (a Access) GetDBSecurityGroups() (*SecurityGroup, error) {
	url := fmt.Sprintf("%s%s/security-groups", RDB_URL, a.TenantID)
	body, err := a.baseRequest(url, "GET", nil)

	type resp struct {
		SecurityGroups []SecurityGroup `json:"security-groups"`
	}
	sr := &resp{}
	err = json.Unmarshal(body, sr)
	if err != nil {
		return nil, err
	}
	return &sr.SecurityGroups, nil
}

type DBInstance struct {
	Created string `json:"created"`
	Id      string `json:"id"`
	Links   []Link `json:"links"`
	Name    string `json:"name"`
	Status  string `json:"name"`
	Flavor  struct {
		Name  string `json:"name"`
		ID    string `json:"id"`
		Links []Link `json:"links"`
	} `json:"flavor"`
}

type DBInstances struct {
	Instances []DBInstance `json:"instances"`
}

type Database struct {
	Name      string `json:"name"`
	FlavorRef string `json:"flavorRef"`
	Port      int    `json:"port"`
	DBType    struct {
		Name    string `json:"name"`
		Version string `json:"version"`
	} `json:"dbtype"`
}

type DBFlavors struct {
	Flavors []DBFlavor `json:"flavors"`
}

/*
 Instance Details type that is returned by server
*/
type InstDetails struct {
	Created        string          `json:"created"`
	Hostname       string          `json:"hostname"`
	ID             string          `json:"id"`
	Links          []Link          `json:"links"`
	Name           string          `json:"name"`
	Port           int             `json:"port"`
	SecurityGroups []SecurityGroup `json:"security_groups"`
	Status         string          `json:"status"`
	Updated        string          `json:"updated"`
	Flavor         struct {
		Name  string `json:"name"`
		ID    string `json:"id"`
		Links []Link `json:"links"`
	} `json:"flavor"`
}

/*
 Type describing database flavor
*/
type DBFlavor struct {
	Id    int    `json:"id"`
	Links []Link `json:"links"`
	Name  string `json:"name"`
	Ram   int    `json:"ram"`
	Vcpu  int    `json:"vcpu"`
}

/*
 This type describes the JSON data which should be sent to the
create database instance resource.
*/
type DatabaseReq struct {
	Instance Database `json:"instance"`
}

/*
 This type describes JSON response from a successful CreateDBInstance
 call.
*/
type NewDBInstance struct {
	Created    string        `json:"created"`
	Credential DBCredentials `json:"credential"`
	Flavor     struct {
		Name  string `json:"name"`
		ID    string `json:"id"`
		Links []Link `json:"links"`
	} `json:"flavor"`
	Hostname       string        `json:"hostname"`
	Id             string        `json:"id"`
	Links          []Link        `json:"links"`
	Name           string        `json:"name"`
	SecurityGroups []DBSecGroups `json:"security_groups"`
	Status         string        `json:"status"`
}

/*
 This type describes Database Security groups
*/
type DBSecGroups struct {
	Id    string `json:"id"`
	Links []Link `json:"links"`
}

/*
 This type describes Database Credentials
*/
type DBCredentials struct {
	Password string `json:"password"`
	Username string `json:"username"`
}

func (f DBFlavors) GetFlavorRef(fn string) string {
	for _, val := range f.Flavors {
		if val.Name == fn {
			return val.Links[0].HREF
		}
	}
	panic("Flavor not found")
}

func (db DatabaseReq) MarshalJSON() ([]byte, error) {
	b := bytes.NewBufferString(`{"instance":{`)
	if db.Instance.Name == "" {
		return nil, errors.New("A name is required")
	}
	b.WriteString(fmt.Sprintf(`"name":"%s",`, db.Instance.Name))
	if db.Instance.FlavorRef == "" {
		return nil, errors.New("Flavor is required")
	}
	b.WriteString(fmt.Sprintf(`"flavorRef":"%s",`,
		db.Instance.FlavorRef))
	if db.Instance.Port == 0 {
		b.WriteString(`"port":"3306",`)
	} else {
		b.WriteString(fmt.Sprintf(`"port":"%s",`, db.Instance.Port))
	}
	b.WriteString(`"dbtype":{`)
	b.WriteString(`"name":"mysql",`)
	b.WriteString(`"version":"5.5"}}}`)

	return b.Bytes(), nil
}
<commit_msg>added security group details<commit_after>// Copyright (c) 2013, Bartosz Kratochwil
// All rights reserved.

// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:

//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.

//     * Redistributions in binary form must reproduce the above
//       copyright notice, this list of conditions and the following
//       disclaimer in the documentation and/or other materials provided
//       with the distribution.

//     * Neither the name of Aaron France nor the names of its
//       contributors may be used to endorse or promote products derived
//       from this software without specific prior written permission.

// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
package hpcloud

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
)

/*
 ListDBInstances will list all the available database instances

This function implements the interface as described in:
http://api-docs.hpcloud.com/hpcloud-rdb-mysql/1.0/content/list-database-instances.html
*/
func (a Access) ListDBInstances() (*DBInstances, error) {
	url := fmt.Sprintf("%s%s/instances", RDB_URL, a.TenantID)
	body, err := a.baseRequest(url, "GET", nil)
	if err != nil {
		return nil, err
	}
	dbs := &DBInstances{}
	err = json.Unmarshal(body, dbs)
	return dbs, err
}

/*
 This function takes instance ID and deletes database instance with this ID.

This function implements the interface as described in:
http://api-docs.hpcloud.com/hpcloud-rdb-mysql/1.0/content/delete-instance.html
*/
func (a Access) DeleteDBInstance(instanceID string) error {
	url := fmt.Sprintf("%s%s/instances/%s", RDB_URL, a.TenantID,
		instanceID)
	_, err := a.baseRequest(url, "DELETE", nil)
	return err
}

/*
 This function takes instance ID and restarts DB instance with this ID.

This function implements the interface as described in:
http://api-docs.hpcloud.com/hpcloud-rdb-mysql/1.0/content/restart-instance.html
*/
func (a Access) RestartDBInstance(instanceID string) error {
	b := `{"restart":{}}`
	url := fmt.Sprintf("%s%s/instances/%s/action", RDB_URL,
		a.TenantID, instanceID)
	_, err := a.baseRequest(url, "POST", strings.NewReader(b))
	return err
}

/*
 ListAllFlavors lists all available database flavors.
 This function implements interface as described in:-
http://api-docs.hpcloud.com/hpcloud-rdb-mysql/1.0/content/list-flavors.html
*/
func (a Access) ListAllFlavors() (*DBFlavors, error) {
	url := fmt.Sprintf("%s%s/flavors", RDB_URL, a.TenantID)
	body, err := a.baseRequest(url, "GET", nil)
	if err != nil {
		return nil, err
	}

	flv := &DBFlavors{}
	err = json.Unmarshal(body, flv)
	return flv, err
}

/*
 This function returns flavor specs for given flavor.

 This function implements the interface as described in:
 http://api-docs.hpcloud.com/hpcloud-rdb-mysql/1.0/content/get-flavor.html
*/
func (a Access) GetDBFlavor(ID string) (*DBFlavor, error) {
	url := fmt.Sprintf("%s/flavors/%s", RDB_URL, ID)
	body, err := a.baseRequest(url, "GET", nil)
	if err != nil {
		return nil, err
	}

	flv := &DBFlavor{}
	err = json.Unmarshal(body, flv)
	if err != nil {
		return nil, err
	}
	return flv, nil
}

/*
 CreateDBInstance creates new database instance in the HPCloud using
settings found in the DatabaseReq instance passed to this function

 This function implements the interface as described in:
 http://api-docs.hpcloud.com/hpcloud-rdb-mysql/1.0/content/create-instance.html
*/
func (a Access) CreateDBInstance(db DatabaseReq) (*NewDBInstance, error) {
	b, err := db.MarshalJSON()
	if err != nil {
		return nil, err
	}

	url := fmt.Sprintf("%s%s/instances", RDB_URL, a.TenantID)

	body, err := a.baseRequest(url, "POST",
		strings.NewReader(string(b)))
	if err != nil {
		return nil, err
	}

	type respDB struct {
		Instance NewDBInstance `json:"instance"`
	}

	sr := &respDB{}
	err = json.Unmarshal(body, sr)
	if err != nil {
		return nil, err
	}
	return &sr.Instance, nil
}

/*
 This function retrieves details of the instance with provided ID.

This function implements the interface as described in:
http://api-docs.hpcloud.com/hpcloud-rdb-mysql/1.0/content/get-instance.html
*/
func (a Access) GetDBInstance(id string) (*InstDetails, error) {
	url := fmt.Sprintf("%s%s/instances/%s", RDB_URL, a.TenantID, id)
	body, err := a.baseRequest(url, "GET", nil)
	if err != nil {
		return nil, err
	}
	det := &InstDetails{}
	err = json.Unmarshal(body, det)
	return det, err
}

/*
 This function takes instance ID and resets password for this instance. It
 returns a new instance password.

 This function implements the interface as decribed in:
 http://api-docs.hpcloud.com/hpcloud-rdb-mysql/1.0/content/reset-instance-password.html
*/
func (a Access) ResetDBPassword(id string) (*DBCredentials, error) {
	b := `{"reset-password":{}}`
	url := fmt.Sprintf("%s%s/instances/%s/action", RDB_URL,
		a.TenantID, id)
	body, err := a.baseRequest(url, "POST", strings.NewReader(b))

	sr := &DBCredentials{}
	err = json.Unmarshal(body, sr)
	if err != nil {
		return nil, err
	}
	return sr, nil
}

/*
 This function lists all the security groups available for tenant.

 This function implements the interface as described in:
 http://api-docs.hpcloud.com/hpcloud-rdb-mysql/1.0/content/list-security-groups.html
*/
func (a Access) GetDBSecurityGroups() (*[]SecurityGroup, error) {
	url := fmt.Sprintf("%s%s/security-groups", RDB_URL, a.TenantID)
	body, err := a.baseRequest(url, "GET", nil)

	type resp struct {
		SecurityGroups []SecurityGroup `json:"security_groups"`
	}
	sr := &resp{}
	err = json.Unmarshal(body, sr)
	if err != nil {
		return nil, err
	}
	return &sr.SecurityGroups, nil
}

/*
 This function lists specific security group.

 This function implements the interface as described in:
 http://api-docs.hpcloud.com/hpcloud-rdb-mysql/1.0/content/get-security-group.html
*/
func (a Access) DBSecGroupDetails(sg string) (*SecurityGroup, error) {
	url := fmt.Sprintf("%s%s/security-groups/%s", RDB_URL, a.TenantID, sg)
	body, err := a.baseRequest(url, "GET", nil)

	type resp struct {
		SecurityGroup SecurityGroup `json:"security_group"`
	}

	sr := &resp{}
	err = json.Unmarshal(body, sr)
	if err != nil {
		return nil, err
	}
	return &sr.SecurityGroup, nil
}

type DBInstance struct {
	Created string `json:"created"`
	Id      string `json:"id"`
	Links   []Link `json:"links"`
	Name    string `json:"name"`
	Status  string `json:"name"`
	Flavor  struct {
		Name  string `json:"name"`
		ID    string `json:"id"`
		Links []Link `json:"links"`
	} `json:"flavor"`
}

type DBInstances struct {
	Instances []DBInstance `json:"instances"`
}

type Database struct {
	Name      string `json:"name"`
	FlavorRef string `json:"flavorRef"`
	Port      int    `json:"port"`
	DBType    struct {
		Name    string `json:"name"`
		Version string `json:"version"`
	} `json:"dbtype"`
}

type DBFlavors struct {
	Flavors []DBFlavor `json:"flavors"`
}

/*
 Instance Details type that is returned by server
*/
type InstDetails struct {
	Created        string          `json:"created"`
	Hostname       string          `json:"hostname"`
	ID             string          `json:"id"`
	Links          []Link          `json:"links"`
	Name           string          `json:"name"`
	Port           int             `json:"port"`
	SecurityGroups []SecurityGroup `json:"security_groups"`
	Status         string          `json:"status"`
	Updated        string          `json:"updated"`
	Flavor         struct {
		Name  string `json:"name"`
		ID    string `json:"id"`
		Links []Link `json:"links"`
	} `json:"flavor"`
}

/*
 Type describing database flavor
*/
type DBFlavor struct {
	Id    int    `json:"id"`
	Links []Link `json:"links"`
	Name  string `json:"name"`
	Ram   int    `json:"ram"`
	Vcpu  int    `json:"vcpu"`
}

/*
 This type describes the JSON data which should be sent to the
create database instance resource.
*/
type DatabaseReq struct {
	Instance Database `json:"instance"`
}

/*
 This type describes JSON response from a successful CreateDBInstance
 call.
*/
type NewDBInstance struct {
	Created    string        `json:"created"`
	Credential DBCredentials `json:"credential"`
	Flavor     struct {
		Name  string `json:"name"`
		ID    string `json:"id"`
		Links []Link `json:"links"`
	} `json:"flavor"`
	Hostname       string        `json:"hostname"`
	Id             string        `json:"id"`
	Links          []Link        `json:"links"`
	Name           string        `json:"name"`
	SecurityGroups []DBSecGroups `json:"security_groups"`
	Status         string        `json:"status"`
}

/*
 This type describes Database Security groups
*/
type DBSecGroups struct {
	Id    string `json:"id"`
	Links []Link `json:"links"`
}

/*
 This type describes Database Credentials
*/
type DBCredentials struct {
	Password string `json:"password"`
	Username string `json:"username"`
}

func (f DBFlavors) GetFlavorRef(fn string) string {
	for _, val := range f.Flavors {
		if val.Name == fn {
			return val.Links[0].HREF
		}
	}
	panic("Flavor not found")
}

func (db DatabaseReq) MarshalJSON() ([]byte, error) {
	b := bytes.NewBufferString(`{"instance":{`)
	if db.Instance.Name == "" {
		return nil, errors.New("A name is required")
	}
	b.WriteString(fmt.Sprintf(`"name":"%s",`, db.Instance.Name))
	if db.Instance.FlavorRef == "" {
		return nil, errors.New("Flavor is required")
	}
	b.WriteString(fmt.Sprintf(`"flavorRef":"%s",`,
		db.Instance.FlavorRef))
	if db.Instance.Port == 0 {
		b.WriteString(`"port":"3306",`)
	} else {
		b.WriteString(fmt.Sprintf(`"port":"%s",`, db.Instance.Port))
	}
	b.WriteString(`"dbtype":{`)
	b.WriteString(`"name":"mysql",`)
	b.WriteString(`"version":"5.5"}}}`)

	return b.Bytes(), nil
}
<|endoftext|>


<commit_before>// Package hub provides a simple event dispatcher for publish/subscribe pattern.
// Optimized for publish calls.
package hub

import "sync"

// Event is an interface for published events.
type Event interface {
	Kind() int
}

// Hub is an event dispatcher, publishes events to the subscribers
// which are subscribed for a specific event type.
type Hub struct {
	subscribers map[int][]handler
	m           sync.RWMutex
	seq         uint64
}

type handler struct {
	f  func(Event)
	id uint64
}

// New returns pointer to a new Hub.
func New() *Hub {
	return &Hub{subscribers: make(map[int][]handler)}
}

// Subscribe registers f for the event of a specific kind.
func (h *Hub) Subscribe(kind int, f func(Event)) (cancel func()) {
	h.m.Lock()
	h.seq++
	id := h.seq
	h.subscribers[kind] = append(h.subscribers[kind], handler{id: id, f: f})
	h.m.Unlock()
	return func() {
		h.m.Lock()
		a := h.subscribers[kind]
		for i, h := range a {
			if h.id == id {
				a[i], a = a[len(a)-1], a[:len(a)-1]
				break
			}
		}
		if len(a) == 0 {
			delete(h.subscribers, kind)
		}
		h.m.Unlock()
	}
}

// Publish an event to the subscribers.
func (h *Hub) Publish(e Event) {
	h.m.RLock()
	if handlers, ok := h.subscribers[e.Kind()]; ok {
		for _, h := range handlers {
			h.f(e)
		}
	}
	h.m.RUnlock()
}

// DefaultHub is the default Hub used by Publish and Subscribe.
var DefaultHub = New()

// Subscribe registers f for the event of a specific kind in the DefaultHub.
func Subscribe(kind int, f func(Event)) (cancel func()) {
	return DefaultHub.Subscribe(kind, f)
}

// Publish an event to the subscribers in DefaultHub.
func Publish(e Event) {
	DefaultHub.Publish(e)
}
<commit_msg>update comments<commit_after>// Package hub provides a simple event dispatcher for publish/subscribe pattern.
package hub

import "sync"

// Event is an interface for published events.
type Event interface {
	Kind() int
}

// Hub is an event dispatcher, publishes events to the subscribers
// which are subscribed for a specific event type.
// Optimized for publish calls.
// The handlers may be called in order different than they are registered.
type Hub struct {
	subscribers map[int][]handler
	m           sync.RWMutex
	seq         uint64
}

type handler struct {
	f  func(Event)
	id uint64
}

// New returns pointer to a new Hub.
func New() *Hub {
	return &Hub{subscribers: make(map[int][]handler)}
}

// Subscribe registers f for the event of a specific kind.
func (h *Hub) Subscribe(kind int, f func(Event)) (cancel func()) {
	h.m.Lock()
	h.seq++
	id := h.seq
	h.subscribers[kind] = append(h.subscribers[kind], handler{id: id, f: f})
	h.m.Unlock()
	return func() {
		h.m.Lock()
		a := h.subscribers[kind]
		for i, h := range a {
			if h.id == id {
				a[i], a = a[len(a)-1], a[:len(a)-1]
				break
			}
		}
		if len(a) == 0 {
			delete(h.subscribers, kind)
		}
		h.m.Unlock()
	}
}

// Publish an event to the subscribers.
func (h *Hub) Publish(e Event) {
	h.m.RLock()
	if handlers, ok := h.subscribers[e.Kind()]; ok {
		for _, h := range handlers {
			h.f(e)
		}
	}
	h.m.RUnlock()
}

// DefaultHub is the default Hub used by Publish and Subscribe.
var DefaultHub = New()

// Subscribe registers f for the event of a specific kind in the DefaultHub.
func Subscribe(kind int, f func(Event)) (cancel func()) {
	return DefaultHub.Subscribe(kind, f)
}

// Publish an event to the subscribers in DefaultHub.
func Publish(e Event) {
	DefaultHub.Publish(e)
}
<|endoftext|>


<commit_before>package net4g

import (
	"sync"
	"time"
	"github.com/carsonsx/log4g"
	"runtime/debug"
	"errors"
	"fmt"
)

type NetHub struct {
	connections      map[string]NetConn
	groupsConn       map[string][]NetConn
	groupsRound      map[string]int
	groupMutex       sync.Mutex
	enableLB         bool
	lbConnections    []NetConn
	lbRound          int
	addChan          chan *chanData
	getChan          chan *chanData
	hasChan          chan *chanData
	sliceChan        chan *chanData
	keyChan          chan *chanData
	removeChan       chan NetConn
	size             int
	heartbeat        bool
	heartbeatTicker  *time.Ticker
	heartbeatTimeout time.Duration
	kickChan         chan func(session NetSession) bool
	broadcastChan    chan *chanData
	closing          chan bool
	closed           bool
	wg               sync.WaitGroup
}

type chanData struct {
	conn    NetConn
	conns    []NetConn
	key     string
	value   interface{}
	one     bool
	group   string
	data    []byte
	filter  func(session NetSession) bool
	once    bool
	errFunc func(err error)
	wg sync.WaitGroup
}

func (hub *NetHub) Start() {
	hub.connections = make(map[string]NetConn)
	hub.groupsConn = make(map[string][]NetConn)
	hub.groupsRound = make(map[string]int)
	hub.addChan = make(chan *chanData, 100)
	hub.keyChan = make(chan *chanData, 100)
	hub.getChan = make(chan *chanData, 1)
	hub.hasChan = make(chan *chanData, 1)
	hub.sliceChan = make(chan *chanData, 1)
	hub.removeChan = make(chan NetConn, 100)
	hub.kickChan = make(chan func(session NetSession) bool, 10)
	hub.broadcastChan = make(chan *chanData, 1000)
	hub.closing = make(chan bool, 1)
	hub.wg.Add(1)
	go func() {
		hub.heartbeatTimeout = NetConfig.HeartbeatFrequency + NetConfig.NetTolerableTime
		hub.heartbeatTicker = time.NewTicker(HEART_BEAT_INTERVAL)
		if !hub.heartbeat {
			hub.heartbeatTicker.Stop()
		}
		for {
			if !hub.do() {
				break
			}
		}
		hub.wg.Done()
	}()
}

func (hub *NetHub) do() (cond bool) {

	defer func() {
		if r := recover(); r != nil {
			log4g.Error("********************* Hub Panic *********************")
			log4g.Error(r)
			log4g.Error(string(debug.Stack()))
			log4g.Error("********************* Hub Panic *********************")
			cond = true
		}
	}()

	select {
	case cData := <-hub.addChan:
		cData.conn.Session().Set(SESSION_CONNECT_KEY, cData.key)
		hub.connections[cData.key] = cData.conn
		if hub.enableLB {
			hub.lbConnections = append(hub.lbConnections, cData.conn)
		}
		hub.size++
		log4g.Debug("connection count: %d", len(hub.connections))
		cData.wg.Done()
	case filter := <-hub.kickChan:
		if filter != nil {
			for _, conn := range hub.connections {
				if filter(conn.Session()) {
					conn.Close()
					hub._delete(conn)
					log4g.Warn("kicked connection %s", conn.RemoteAddr().String())
					break
				}
			}
		}
	case cData := <-hub.keyChan:
		if _, ok := hub.connections[cData.key]; ok {
			log4g.Panic("connection name '%s' existed", cData.key)
		}
		delete(hub.connections, cData.conn.Session().GetString(SESSION_CONNECT_KEY))
		cData.conn.Session().Set(SESSION_CONNECT_KEY, cData.key)
		hub.connections[cData.key] = cData.conn
		log4g.Debug("connection count: %d", len(hub.connections))
	case cData := <-hub.getChan:
		cData.conn = hub.connections[cData.key]
		cData.wg.Done()
	case cData := <-hub.hasChan:
		cData.one = false
		for _, conn := range hub.connections {
			if conn.Session().Get(cData.key) == cData.value {
				cData.one = true
				break
			}
		}
		cData.wg.Done()
	case cData := <-hub.sliceChan:
		for _, conn := range hub.connections {
			cData.conns = append(cData.conns, conn)
		}
		cData.wg.Done()
	case conn := <-hub.removeChan:
		hub._delete(conn)
	case t := <-hub.heartbeatTicker.C:
		for _, conn := range hub.connections {
			if t.UnixNano() > conn.Session().GetInt64(HEART_BEAT_LAST_TIME)+hub.heartbeatTimeout.Nanoseconds() {
				log4g.Warn("client timeout: %s", conn.RemoteAddr().String())
				hub._delete(conn)
				conn.Close()
			}
		}
	case cData := <-hub.broadcastChan:
		if cData.key != "" { // Send to one conn
			hub.connections[cData.key].Write(cData.data)
		} else if cData.one {
			connCount := len(hub.lbConnections)
			if connCount > 0 {
				if hub.lbRound >= connCount {
					hub.lbRound = 0
				}
				log4g.Trace("group size: %d, round robin: %d", connCount, hub.lbRound)
				err := hub.lbConnections[hub.lbRound].Write(cData.data)
				if err != nil &&cData.errFunc != nil {
					cData.errFunc(err)
				}
				hub.lbRound++
			} else if cData.errFunc != nil {
				err := errors.New("not found any connection")
				log4g.Error(err)
				cData.errFunc(err)
			}
		} else if cData.group != "" {// Send to group
			log4g.Debug("send to group %s", cData.group)
			conns := hub.groupsConn[cData.group]
			if cData.once { // send to group's one by round robin
				round := hub.groupsRound[cData.group]
				groupsConnCount := len(conns)
				if groupsConnCount > 0 {
					if round >= groupsConnCount {
						round = 0
					}
					//log4g.Debug("group size: %d, round robin: %d", groupsConnCount, round)
					conns[round].Write(cData.data)
					round++
					hub.groupsRound[cData.group] = round
				} else {
					err := errors.New(fmt.Sprintf("not found any group[%s] connection", cData.group))
					log4g.Error(err)
					if cData.errFunc != nil {
						cData.errFunc(err)
					}
				}
			} else { // send to group's all
				for _, conn := range conns {
					conn.Write(cData.data)
				}
			}
		} else {// broadcast by filter
			//log4g.Debug("broadcast to %d connections", len(hub.connections))
			if log4g.IsDebugEnabled() {
				log4g.Debug("broadcast msg to %d connections", len(hub.connections))
			}
			for _, conn := range hub.connections {
				if cData.filter == nil || cData.filter(conn.Session()) {
					conn.Write(cData.data)
					if log4g.IsDebugEnabled() {
						log4g.Debug("sent to %s", conn.Session().Get(SESSION_ID))
					}
					if cData.once {
						break
					}
				}
			}
		}
	case <-hub.closing:
		return false
	}

	return true
}

func (hub *NetHub) _delete(conn NetConn) {
	delete(hub.connections, conn.Session().GetString(SESSION_CONNECT_KEY))
	for i, _conn := range hub.lbConnections {
		if _conn == conn {
			hub.lbConnections = append(hub.lbConnections[:i], hub.lbConnections[i+1:]...)
			break
		}
	}
	gName := conn.Session().GetString(SESSION_GROUP_NAME)
	if gName != "" {
		hub.groupMutex.Lock()
		groups := hub.groupsConn[gName]
		for i, c := range groups {
			if c == conn {
				hub.groupsConn[gName] = append(groups[:i], groups[i+1:]...)
				break
			}
		}
		hub.groupMutex.Unlock()
	}
	hub.size--
}

func (hub *NetHub) SetKey(conn NetConn, key string) {
	cData := new(chanData)
	cData.conn = conn
	cData.key = key
	hub.keyChan <- cData
}

func (hub *NetHub) Add(key string, conn NetConn) {
	hub.Heartbeat(conn)
	cData := new(chanData)
	cData.key = key
	cData.conn = conn
	cData.wg.Add(1)
	hub.addChan <- cData
	cData.wg.Wait()
}

func (hub *NetHub) Get(key string) NetConn {
	cData := new(chanData)
	cData.key = key
	cData.wg.Add(1)
	hub.getChan <- cData
	cData.wg.Wait()
	return cData.conn
}

func (hub *NetHub) HasSession(k string, v interface{}) bool {
	log4g.Trace("has session: key=%s,val=%s", k, v)
	cData := new(chanData)
	cData.wg.Add(1)
	cData.key = k
	cData.value = v
	hub.hasChan <- cData
	cData.wg.Wait()
	log4g.Trace("has session: %v",  cData.one)
	return cData.one
}

func (hub *NetHub) Remove(conn NetConn) {
	hub.removeChan <- conn
}

func (hub *NetHub) Slice() []NetConn {
	cData := new(chanData)
	cData.wg.Add(1)
	hub.sliceChan <- cData
	cData.wg.Wait()
	return cData.conns
}

func (hub *NetHub) Heartbeat(conn NetConn) {
	if hub.heartbeat {
		conn.Session().Set(HEART_BEAT_LAST_TIME, time.Now().UnixNano())
	}
}

func (hub *NetHub) Kick(filter func(session NetSession) bool) {
	hub.kickChan <- filter
}

func (hub *NetHub) Broadcast(data []byte, filter func(session NetSession) bool) {
	cData := new(chanData)
	cData.data = data
	cData.filter = filter
	hub.broadcastChan <- cData
}

func (hub *NetHub) BroadcastAll(data []byte) {
	hub.Broadcast(data, nil)
}

func (hub *NetHub) BroadcastOthers(mySession NetSession, data []byte) {
	hub.Broadcast(data, func(session NetSession) bool {
		return mySession != session
	})
}

func (hub *NetHub) Someone(data []byte, filter func(session NetSession) bool) {
	cData := new(chanData)
	cData.data = data
	cData.filter = filter
	cData.once = true
	hub.broadcastChan <- cData
}

func (hub *NetHub) SetGroup(session NetSession, group string) {
	if group != "" {
		hub.groupMutex.Lock()
		conn := hub.Get(session.GetString(SESSION_CONNECT_KEY))
		if conn != nil {
			groups := hub.groupsConn[group]
			hub.groupsConn[group] = append(groups, conn)
			session.Set(SESSION_GROUP_NAME, group)
			log4g.Debug("set group %s for %s", group, conn.RemoteAddr().String())
			log4g.Debug("group size: %d", len(hub.groupsConn[group]))
		}
		hub.groupMutex.Unlock()
	}
}

func (hub *NetHub) Group(group string, data []byte) {
	cData := new(chanData)
	cData.group = group
	cData.data = data
	hub.broadcastChan <- cData
}

func (hub *NetHub) GroupOne(group string, data []byte, errFunc func(error)) {
	cData := new(chanData)
	cData.group = group
	cData.data = data
	cData.once = true
	cData.errFunc = errFunc
	hub.broadcastChan <- cData
}

func (hub *NetHub) One(data []byte, errFunc func(error)) error {

	if !hub.enableLB {
		panic("required to enable load balance")
	}

	if hub.closed {
		text := "hub was closed"
		log4g.Error(text)
		return errors.New(text)
	}

	cData := new(chanData)
	cData.one = true
	cData.data = data
	cData.errFunc = errFunc
	hub.broadcastChan <- cData
	return nil
}

func (hub *NetHub) CloseConnections() {
	for _, conn := range hub.connections {
		conn.Close()
	}
	log4g.Debug("closed %d connections", len(hub.connections))
}

func (hub *NetHub) Closed() bool {
	return hub.closed
}

func (hub *NetHub) Destroy() {
	hub.closing <- true
	hub.wg.Wait()
	//close(hub.addChan)
	//close(hub.removeChan)
	//close(hub.broadcastChan)
	//close(hub.closing)
	hub.closed = true
	log4g.Info("closed net hub")
}
<commit_msg>update log<commit_after>package net4g

import (
	"sync"
	"time"
	"github.com/carsonsx/log4g"
	"runtime/debug"
	"errors"
	"fmt"
)

type NetHub struct {
	connections      map[string]NetConn
	groupsConn       map[string][]NetConn
	groupsRound      map[string]int
	groupMutex       sync.Mutex
	enableLB         bool
	lbConnections    []NetConn
	lbRound          int
	addChan          chan *chanData
	getChan          chan *chanData
	hasChan          chan *chanData
	sliceChan        chan *chanData
	keyChan          chan *chanData
	removeChan       chan NetConn
	size             int
	heartbeat        bool
	heartbeatTicker  *time.Ticker
	heartbeatTimeout time.Duration
	kickChan         chan func(session NetSession) bool
	broadcastChan    chan *chanData
	closing          chan bool
	closed           bool
	wg               sync.WaitGroup
}

type chanData struct {
	conn    NetConn
	conns    []NetConn
	key     string
	value   interface{}
	one     bool
	group   string
	data    []byte
	filter  func(session NetSession) bool
	once    bool
	errFunc func(err error)
	wg sync.WaitGroup
}

func (hub *NetHub) Start() {
	hub.connections = make(map[string]NetConn)
	hub.groupsConn = make(map[string][]NetConn)
	hub.groupsRound = make(map[string]int)
	hub.addChan = make(chan *chanData, 100)
	hub.keyChan = make(chan *chanData, 100)
	hub.getChan = make(chan *chanData, 1)
	hub.hasChan = make(chan *chanData, 1)
	hub.sliceChan = make(chan *chanData, 1)
	hub.removeChan = make(chan NetConn, 100)
	hub.kickChan = make(chan func(session NetSession) bool, 10)
	hub.broadcastChan = make(chan *chanData, 1000)
	hub.closing = make(chan bool, 1)
	hub.wg.Add(1)
	go func() {
		hub.heartbeatTimeout = NetConfig.HeartbeatFrequency + NetConfig.NetTolerableTime
		hub.heartbeatTicker = time.NewTicker(HEART_BEAT_INTERVAL)
		if !hub.heartbeat {
			hub.heartbeatTicker.Stop()
		}
		for {
			if !hub.do() {
				break
			}
		}
		hub.wg.Done()
	}()
}

func (hub *NetHub) do() (cond bool) {

	defer func() {
		if r := recover(); r != nil {
			log4g.Error("********************* Hub Panic *********************")
			log4g.Error(r)
			log4g.Error(string(debug.Stack()))
			log4g.Error("********************* Hub Panic *********************")
			cond = true
		}
	}()

	select {
	case cData := <-hub.addChan:
		cData.conn.Session().Set(SESSION_CONNECT_KEY, cData.key)
		hub.connections[cData.key] = cData.conn
		if hub.enableLB {
			hub.lbConnections = append(hub.lbConnections, cData.conn)
		}
		hub.size++
		log4g.Debug("connection count: %d", len(hub.connections))
		cData.wg.Done()
	case filter := <-hub.kickChan:
		if filter != nil {
			for _, conn := range hub.connections {
				if filter(conn.Session()) {
					conn.Close()
					hub._delete(conn)
					log4g.Warn("kicked connection %s", conn.RemoteAddr().String())
					break
				}
			}
		}
	case cData := <-hub.keyChan:
		if _, ok := hub.connections[cData.key]; ok {
			log4g.Panic("connection name '%s' existed", cData.key)
		}
		delete(hub.connections, cData.conn.Session().GetString(SESSION_CONNECT_KEY))
		cData.conn.Session().Set(SESSION_CONNECT_KEY, cData.key)
		hub.connections[cData.key] = cData.conn
		log4g.Debug("connection count: %d", len(hub.connections))
	case cData := <-hub.getChan:
		cData.conn = hub.connections[cData.key]
		cData.wg.Done()
	case cData := <-hub.hasChan:
		cData.one = false
		for _, conn := range hub.connections {
			if conn.Session().Get(cData.key) == cData.value {
				cData.one = true
				break
			}
		}
		cData.wg.Done()
	case cData := <-hub.sliceChan:
		for _, conn := range hub.connections {
			cData.conns = append(cData.conns, conn)
		}
		cData.wg.Done()
	case conn := <-hub.removeChan:
		hub._delete(conn)
	case t := <-hub.heartbeatTicker.C:
		for _, conn := range hub.connections {
			if t.UnixNano() > conn.Session().GetInt64(HEART_BEAT_LAST_TIME)+hub.heartbeatTimeout.Nanoseconds() {
				log4g.Warn("client timeout: %s", conn.RemoteAddr().String())
				hub._delete(conn)
				conn.Close()
			}
		}
	case cData := <-hub.broadcastChan:
		if cData.key != "" { // Send to one conn
			hub.connections[cData.key].Write(cData.data)
		} else if cData.one {
			connCount := len(hub.lbConnections)
			if connCount > 0 {
				if hub.lbRound >= connCount {
					hub.lbRound = 0
				}
				log4g.Trace("group size: %d, round robin: %d", connCount, hub.lbRound)
				err := hub.lbConnections[hub.lbRound].Write(cData.data)
				if err != nil &&cData.errFunc != nil {
					cData.errFunc(err)
				}
				hub.lbRound++
			} else if cData.errFunc != nil {
				err := errors.New("not found any connection")
				log4g.Error(err)
				cData.errFunc(err)
			}
		} else if cData.group != "" {// Send to group
			log4g.Debug("send to group %s", cData.group)
			conns := hub.groupsConn[cData.group]
			if cData.once { // send to group's one by round robin
				round := hub.groupsRound[cData.group]
				groupsConnCount := len(conns)
				if groupsConnCount > 0 {
					if round >= groupsConnCount {
						round = 0
					}
					//log4g.Debug("group size: %d, round robin: %d", groupsConnCount, round)
					conns[round].Write(cData.data)
					round++
					hub.groupsRound[cData.group] = round
				} else {
					err := errors.New(fmt.Sprintf("not found any group[%s] connection", cData.group))
					log4g.Error(err)
					if cData.errFunc != nil {
						cData.errFunc(err)
					}
				}
			} else { // send to group's all
				for _, conn := range conns {
					conn.Write(cData.data)
				}
			}
		} else {// broadcast by filter
			//log4g.Debug("broadcast to %d connections", len(hub.connections))
			for _, conn := range hub.connections {
				if cData.filter == nil || cData.filter(conn.Session()) {
					conn.Write(cData.data)
					if log4g.IsDebugEnabled() {
						log4g.Debug("[broadcast] sent to %s", conn.Session().Get(SESSION_ID))
					}
					if cData.once {
						break
					}
				}
			}
		}
	case <-hub.closing:
		return false
	}

	return true
}

func (hub *NetHub) _delete(conn NetConn) {
	delete(hub.connections, conn.Session().GetString(SESSION_CONNECT_KEY))
	for i, _conn := range hub.lbConnections {
		if _conn == conn {
			hub.lbConnections = append(hub.lbConnections[:i], hub.lbConnections[i+1:]...)
			break
		}
	}
	gName := conn.Session().GetString(SESSION_GROUP_NAME)
	if gName != "" {
		hub.groupMutex.Lock()
		groups := hub.groupsConn[gName]
		for i, c := range groups {
			if c == conn {
				hub.groupsConn[gName] = append(groups[:i], groups[i+1:]...)
				break
			}
		}
		hub.groupMutex.Unlock()
	}
	hub.size--
}

func (hub *NetHub) SetKey(conn NetConn, key string) {
	cData := new(chanData)
	cData.conn = conn
	cData.key = key
	hub.keyChan <- cData
}

func (hub *NetHub) Add(key string, conn NetConn) {
	hub.Heartbeat(conn)
	cData := new(chanData)
	cData.key = key
	cData.conn = conn
	cData.wg.Add(1)
	hub.addChan <- cData
	cData.wg.Wait()
}

func (hub *NetHub) Get(key string) NetConn {
	cData := new(chanData)
	cData.key = key
	cData.wg.Add(1)
	hub.getChan <- cData
	cData.wg.Wait()
	return cData.conn
}

func (hub *NetHub) HasSession(k string, v interface{}) bool {
	log4g.Trace("has session: key=%s,val=%s", k, v)
	cData := new(chanData)
	cData.wg.Add(1)
	cData.key = k
	cData.value = v
	hub.hasChan <- cData
	cData.wg.Wait()
	log4g.Trace("has session: %v",  cData.one)
	return cData.one
}

func (hub *NetHub) Remove(conn NetConn) {
	hub.removeChan <- conn
}

func (hub *NetHub) Slice() []NetConn {
	cData := new(chanData)
	cData.wg.Add(1)
	hub.sliceChan <- cData
	cData.wg.Wait()
	return cData.conns
}

func (hub *NetHub) Heartbeat(conn NetConn) {
	if hub.heartbeat {
		conn.Session().Set(HEART_BEAT_LAST_TIME, time.Now().UnixNano())
	}
}

func (hub *NetHub) Kick(filter func(session NetSession) bool) {
	hub.kickChan <- filter
}

func (hub *NetHub) Broadcast(data []byte, filter func(session NetSession) bool) {
	cData := new(chanData)
	cData.data = data
	cData.filter = filter
	hub.broadcastChan <- cData
}

func (hub *NetHub) BroadcastAll(data []byte) {
	hub.Broadcast(data, nil)
}

func (hub *NetHub) BroadcastOthers(mySession NetSession, data []byte) {
	hub.Broadcast(data, func(session NetSession) bool {
		return mySession != session
	})
}

func (hub *NetHub) Someone(data []byte, filter func(session NetSession) bool) {
	cData := new(chanData)
	cData.data = data
	cData.filter = filter
	cData.once = true
	hub.broadcastChan <- cData
}

func (hub *NetHub) SetGroup(session NetSession, group string) {
	if group != "" {
		hub.groupMutex.Lock()
		conn := hub.Get(session.GetString(SESSION_CONNECT_KEY))
		if conn != nil {
			groups := hub.groupsConn[group]
			hub.groupsConn[group] = append(groups, conn)
			session.Set(SESSION_GROUP_NAME, group)
			log4g.Debug("set group %s for %s", group, conn.RemoteAddr().String())
			log4g.Debug("group size: %d", len(hub.groupsConn[group]))
		}
		hub.groupMutex.Unlock()
	}
}

func (hub *NetHub) Group(group string, data []byte) {
	cData := new(chanData)
	cData.group = group
	cData.data = data
	hub.broadcastChan <- cData
}

func (hub *NetHub) GroupOne(group string, data []byte, errFunc func(error)) {
	cData := new(chanData)
	cData.group = group
	cData.data = data
	cData.once = true
	cData.errFunc = errFunc
	hub.broadcastChan <- cData
}

func (hub *NetHub) One(data []byte, errFunc func(error)) error {

	if !hub.enableLB {
		panic("required to enable load balance")
	}

	if hub.closed {
		text := "hub was closed"
		log4g.Error(text)
		return errors.New(text)
	}

	cData := new(chanData)
	cData.one = true
	cData.data = data
	cData.errFunc = errFunc
	hub.broadcastChan <- cData
	return nil
}

func (hub *NetHub) CloseConnections() {
	for _, conn := range hub.connections {
		conn.Close()
	}
	log4g.Debug("closed %d connections", len(hub.connections))
}

func (hub *NetHub) Closed() bool {
	return hub.closed
}

func (hub *NetHub) Destroy() {
	hub.closing <- true
	hub.wg.Wait()
	//close(hub.addChan)
	//close(hub.removeChan)
	//close(hub.broadcastChan)
	//close(hub.closing)
	hub.closed = true
	log4g.Info("closed net hub")
}
<|endoftext|>


<commit_before>/*
Copyright 2021 The TestGrid Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package v1 (api/v1) is the first versioned implementation of the API
package v1

import (
	"errors"
	"fmt"
	"net/http"

	"github.com/gorilla/mux"
	"github.com/sirupsen/logrus"

	"github.com/GoogleCloudPlatform/testgrid/config"
	apipb "github.com/GoogleCloudPlatform/testgrid/pb/api/v1"
	configpb "github.com/GoogleCloudPlatform/testgrid/pb/config"
	"github.com/GoogleCloudPlatform/testgrid/util/gcs"
)

func (s Server) configPath(r *http.Request) (path *gcs.Path, isDefault bool, err error) {
	scope := r.URL.Query().Get("scope")
	if scope != "" {
		path, err = gcs.NewPath(fmt.Sprintf("%s/%s", scope, "config"))
		isDefault = false
		return
	}
	if s.DefaultBucket != "" {
		path, err = gcs.NewPath(fmt.Sprintf("%s/%s", s.DefaultBucket, "config"))
		isDefault = true
		return
	}
	return nil, false, errors.New("no testgrid scope")
}

// passQueryParameters returns only the query parameters in the request that need to be passed through to links
func passQueryParameters(r *http.Request) string {
	if scope := r.URL.Query().Get("scope"); scope != "" {
		return fmt.Sprintf("?scope=%s", scope)
	}
	return ""
}

// getConfig will return a config file or will send an error to the http writer
// If this function returns nil, no further writes should be made to 'w'
func (s Server) getConfig(w http.ResponseWriter, r *http.Request) *configpb.Configuration {
	configPath, isDefault, err := s.configPath(r)
	if err != nil || configPath == nil {
		http.Error(w, "Scope not specified", http.StatusBadRequest)
		return nil
	}

	cfg, err := config.ReadGCS(r.Context(), s.Client, *configPath)
	if err != nil {
		http.Error(w, fmt.Sprintf("Could not read config at %q", configPath.String()), http.StatusInternalServerError)
		// Only log an error if we set and use a default scope, but can't access it.
		// Otherwise, invalid requests will write useless logs.
		if isDefault {
			// TODO(chases2): Pass a logrus logger through the server object
			logrus.WithError(err).Errorf("Can't read default config at %q; check permissions", configPath.String())
		}
		return nil
	}
	return cfg
}

// ListDashboardGroups returns every dashboard group in TestGrid
// Response Proto: ListDashboardGroupResponse
func (s Server) ListDashboardGroups(w http.ResponseWriter, r *http.Request) {
	cfg := s.getConfig(w, r)
	if cfg == nil {
		return
	}

	var groups apipb.ListDashboardGroupResponse
	for _, group := range cfg.DashboardGroups {
		rsc := apipb.Resource{
			Name: group.Name,
			Link: fmt.Sprintf("%s/dashboard-groups/%s%s", s.Host, config.Normalize(group.Name), passQueryParameters(r)),
		}
		groups.DashboardGroups = append(groups.DashboardGroups, &rsc)
	}

	writeJSON(w, &groups)
}

// GetDashboardGroup returns a given dashboard group
// Response Proto: GetDashboardGroupResponse
func (s Server) GetDashboardGroup(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	cfg := s.getConfig(w, r)
	if cfg == nil {
		return
	}

	dashboardGroupKey := vars["dashboard-group"]
	for _, group := range cfg.DashboardGroups {
		if config.Normalize(group.Name) == dashboardGroupKey {
			result := apipb.GetDashboardGroupResponse{}
			for _, dash := range group.DashboardNames {
				rsc := apipb.Resource{
					Name: dash,
					Link: fmt.Sprintf("%s/dashboards/%s%s", s.Host, config.Normalize(dash), passQueryParameters(r)),
				}
				result.Dashboards = append(result.Dashboards, &rsc)
			}
			writeJSON(w, &result)
			return
		}
	}
	http.Error(w, fmt.Sprintf("Dashboard group %q not found", vars["dashboard-group"]), http.StatusNotFound)
}

// ListDashboards returns every dashboard in TestGrid
// Response Proto: ListDashboardResponse
func (s Server) ListDashboards(w http.ResponseWriter, r *http.Request) {
	cfg := s.getConfig(w, r)
	if cfg == nil {
		return
	}

	var dashboardResponse apipb.ListDashboardResponse
	for _, dashboard := range cfg.Dashboards {
		rsc := apipb.Resource{
			Name: dashboard.Name,
			Link: fmt.Sprintf("%s/dashboards/%s%s", s.Host, config.Normalize(dashboard.Name), passQueryParameters(r)),
		}
		dashboardResponse.Dashboards = append(dashboardResponse.Dashboards, &rsc)
	}

	writeJSON(w, &dashboardResponse)
}

// GetDashboard returns a given dashboard
// Response Proto: GetDashboardResponse
func (s Server) GetDashboard(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	cfg := s.getConfig(w, r)
	if cfg == nil {
		return
	}

	dashboardKey := vars["dashboard"]
	for _, dashboard := range cfg.Dashboards {
		if config.Normalize(dashboard.Name) == dashboardKey {
			result := apipb.GetDashboardResponse{
				DefaultTab:          dashboard.DefaultTab,
				HighlightToday:      dashboard.HighlightToday,
				SuppressFailingTabs: dashboard.DownplayFailingTabs,
				Notifications:       dashboard.Notifications,
			}
			writeJSON(w, &result)
			return
		}
	}
	http.Error(w, fmt.Sprintf("Dashboard %q not found", vars["dashboard"]), http.StatusNotFound)
}

// ListDashboardTabs returns a given dashboard tabs
// Response Proto: ListDashboardTabsResponse
func (s Server) ListDashboardTabs(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	cfg := s.getConfig(w, r)
	if cfg == nil {
		return
	}

	dashboardKey := vars["dashboard"]
	var dashboardTabsResponse apipb.ListDashboardTabsResponse
	for _, dashboard := range cfg.Dashboards {
		if config.Normalize(dashboard.Name) == dashboardKey {
			for _, tab := range dashboard.DashboardTab {
				rsc := apipb.Resource{
					Name: tab.Name,
					Link: fmt.Sprintf("%s/dashboards/%s/tabs/%s%s", s.Host, config.Normalize(dashboard.Name), config.Normalize(tab.Name), passQueryParameters(r)),
				}
				dashboardTabsResponse.DashboardTabs = append(dashboardTabsResponse.DashboardTabs, &rsc)
			}
			writeJSON(w, &dashboardTabsResponse)
			return
		}
	}
	http.Error(w, fmt.Sprintf("Dashboard %q not found", vars["dashboard"]), http.StatusNotFound)
}
<commit_msg>Update config APIs response type<commit_after>/*
Copyright 2021 The TestGrid Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package v1 (api/v1) is the first versioned implementation of the API
package v1

import (
	"errors"
	"fmt"
	"net/http"

	"github.com/gorilla/mux"
	"github.com/sirupsen/logrus"

	"github.com/GoogleCloudPlatform/testgrid/config"
	apipb "github.com/GoogleCloudPlatform/testgrid/pb/api/v1"
	configpb "github.com/GoogleCloudPlatform/testgrid/pb/config"
	"github.com/GoogleCloudPlatform/testgrid/util/gcs"
)

func (s Server) configPath(r *http.Request) (path *gcs.Path, isDefault bool, err error) {
	scope := r.URL.Query().Get("scope")
	if scope != "" {
		path, err = gcs.NewPath(fmt.Sprintf("%s/%s", scope, "config"))
		isDefault = false
		return
	}
	if s.DefaultBucket != "" {
		path, err = gcs.NewPath(fmt.Sprintf("%s/%s", s.DefaultBucket, "config"))
		isDefault = true
		return
	}
	return nil, false, errors.New("no testgrid scope")
}

// passQueryParameters returns only the query parameters in the request that need to be passed through to links
func passQueryParameters(r *http.Request) string {
	if scope := r.URL.Query().Get("scope"); scope != "" {
		return fmt.Sprintf("?scope=%s", scope)
	}
	return ""
}

// getConfig will return a config file or will send an error to the http writer
// If this function returns nil, no further writes should be made to 'w'
func (s Server) getConfig(w http.ResponseWriter, r *http.Request) *configpb.Configuration {
	configPath, isDefault, err := s.configPath(r)
	if err != nil || configPath == nil {
		http.Error(w, "Scope not specified", http.StatusBadRequest)
		return nil
	}

	cfg, err := config.ReadGCS(r.Context(), s.Client, *configPath)
	if err != nil {
		http.Error(w, fmt.Sprintf("Could not read config at %q", configPath.String()), http.StatusInternalServerError)
		// Only log an error if we set and use a default scope, but can't access it.
		// Otherwise, invalid requests will write useless logs.
		if isDefault {
			// TODO(chases2): Pass a logrus logger through the server object
			logrus.WithError(err).Errorf("Can't read default config at %q; check permissions", configPath.String())
		}
		return nil
	}
	return cfg
}

// ListDashboardGroups returns every dashboard group in TestGrid
// Response json: ListDashboardGroupResponse
func (s Server) ListDashboardGroups(w http.ResponseWriter, r *http.Request) {
	cfg := s.getConfig(w, r)
	if cfg == nil {
		return
	}

	var groups apipb.ListDashboardGroupResponse
	for _, group := range cfg.DashboardGroups {
		rsc := apipb.Resource{
			Name: group.Name,
			Link: fmt.Sprintf("%s/dashboard-groups/%s%s", s.Host, config.Normalize(group.Name), passQueryParameters(r)),
		}
		groups.DashboardGroups = append(groups.DashboardGroups, &rsc)
	}

	writeJSON(w, &groups)
}

// GetDashboardGroup returns a given dashboard group
// Response json: GetDashboardGroupResponse
func (s Server) GetDashboardGroup(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	cfg := s.getConfig(w, r)
	if cfg == nil {
		return
	}

	dashboardGroupKey := vars["dashboard-group"]
	for _, group := range cfg.DashboardGroups {
		if config.Normalize(group.Name) == dashboardGroupKey {
			result := apipb.GetDashboardGroupResponse{}
			for _, dash := range group.DashboardNames {
				rsc := apipb.Resource{
					Name: dash,
					Link: fmt.Sprintf("%s/dashboards/%s%s", s.Host, config.Normalize(dash), passQueryParameters(r)),
				}
				result.Dashboards = append(result.Dashboards, &rsc)
			}
			writeJSON(w, &result)
			return
		}
	}
	http.Error(w, fmt.Sprintf("Dashboard group %q not found", vars["dashboard-group"]), http.StatusNotFound)
}

// ListDashboards returns every dashboard in TestGrid
// Response json: ListDashboardResponse
func (s Server) ListDashboards(w http.ResponseWriter, r *http.Request) {
	cfg := s.getConfig(w, r)
	if cfg == nil {
		return
	}

	var dashboardResponse apipb.ListDashboardResponse
	for _, dashboard := range cfg.Dashboards {
		rsc := apipb.Resource{
			Name: dashboard.Name,
			Link: fmt.Sprintf("%s/dashboards/%s%s", s.Host, config.Normalize(dashboard.Name), passQueryParameters(r)),
		}
		dashboardResponse.Dashboards = append(dashboardResponse.Dashboards, &rsc)
	}

	writeJSON(w, &dashboardResponse)
}

// GetDashboard returns a given dashboard
// Response json: GetDashboardResponse
func (s Server) GetDashboard(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	cfg := s.getConfig(w, r)
	if cfg == nil {
		return
	}

	dashboardKey := vars["dashboard"]
	for _, dashboard := range cfg.Dashboards {
		if config.Normalize(dashboard.Name) == dashboardKey {
			result := apipb.GetDashboardResponse{
				DefaultTab:          dashboard.DefaultTab,
				HighlightToday:      dashboard.HighlightToday,
				SuppressFailingTabs: dashboard.DownplayFailingTabs,
				Notifications:       dashboard.Notifications,
			}
			writeJSON(w, &result)
			return
		}
	}
	http.Error(w, fmt.Sprintf("Dashboard %q not found", vars["dashboard"]), http.StatusNotFound)
}

// ListDashboardTabs returns a given dashboard tabs
// Response json: ListDashboardTabsResponse
func (s Server) ListDashboardTabs(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	cfg := s.getConfig(w, r)
	if cfg == nil {
		return
	}

	dashboardKey := vars["dashboard"]
	var dashboardTabsResponse apipb.ListDashboardTabsResponse
	for _, dashboard := range cfg.Dashboards {
		if config.Normalize(dashboard.Name) == dashboardKey {
			for _, tab := range dashboard.DashboardTab {
				rsc := apipb.Resource{
					Name: tab.Name,
					Link: fmt.Sprintf("%s/dashboards/%s/tabs/%s%s", s.Host, config.Normalize(dashboard.Name), config.Normalize(tab.Name), passQueryParameters(r)),
				}
				dashboardTabsResponse.DashboardTabs = append(dashboardTabsResponse.DashboardTabs, &rsc)
			}
			writeJSON(w, &dashboardTabsResponse)
			return
		}
	}
	http.Error(w, fmt.Sprintf("Dashboard %q not found", vars["dashboard"]), http.StatusNotFound)
}
<|endoftext|>


<commit_before>package flow

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/materials-commons/mcstore/pkg/app"
)

const ChunkPerms = 0700

// A FlowRequest encapsulates the flowjs protocol for uploading a file. The
// protocol supports extensions to the protocol. We extend the protocol to
// include Materials Commons specific information. It is also expected that
// the data sent by flow or another client will be placed in chunkData.
type Request struct {
	FlowChunkNumber  int32  `json:"flowChunkNumber"`  // The chunk being sent.
	FlowTotalChunks  int32  `json:"flowTotalChunks"`  // The total number of chunks to send.
	FlowChunkSize    int32  `json:"flowChunkSize"`    // The size of the chunk.
	FlowTotalSize    int64  `json:"flowTotalSize"`    // The size of the file being uploaded.
	FlowIdentifier   string `json:"flowIdentifier"`   // A unique identifier used by Flow. Not guaranteed to be a GUID.
	FlowFileName     string `json:"flowFilename"`     // The file name being uploaded.
	FlowRelativePath string `json:"flowRelativePath"` // When available the relative file path.
	ProjectID        string `json:"projectID"`        // Materials Commons Project ID.
	DirectoryID      string `json:"directoryID"`      // Materials Commons Directory ID.
	FileID           string `json:"fileID"`           // Materials Commons File ID.
	Chunk            []byte `json:"-"`                // The file data.
	ChunkHash        string `json:"chunkHash"`        // The computed MD5 hash for the chunk (optional).
	FileHash         string `json:"fileHash"`         // The computed MD5 hash for the file (optional)
}

func (r *Request) UploadID() string {
	return fmt.Sprintf("%s-%s-%s", r.ProjectID, r.DirectoryID, r.FileID)
}

func (r *Request) Write() error {
	path := r.Path()
	err := os.MkdirAll(path, ChunkPerms)
	switch {
	case err != nil:
		return err
	case !r.canWriteTo(path):
		return app.ErrInvalid
	default:
		return ioutil.WriteFile(path, r.Chunk, ChunkPerms)
	}
}

// canWriteTo determines if a particular chunk can be written.
// If the size of the ondisk chunk is smaller than the request
// chunk then that chunk is incomplete and we allow a write to it.
func (r *Request) canWriteTo(path string) bool {
	finfo, err := os.Stat(path)
	switch {
	case err != nil:
		return os.IsNotExist(err)
	case finfo.Size() < int64(r.FlowChunkSize):
		return true
	default:
		return false
	}
}

func (r *Request) Dir() string {
	mcdir := app.MCDir.Path()
	uploadPath := filepath.Join(mcdir, "upload", r.ProjectID, r.DirectoryID, r.FileID)
	return uploadPath
}

func (r *Request) Path() string {
	return filepath.Join(r.Dir(), fmt.Sprintf("%d", r.FlowChunkNumber))
}
<commit_msg>Handle case where block is already uploaded.<commit_after>package flow

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/materials-commons/mcstore/pkg/app"
)

const ChunkPerms = 0700

// A FlowRequest encapsulates the flowjs protocol for uploading a file. The
// protocol supports extensions to the protocol. We extend the protocol to
// include Materials Commons specific information. It is also expected that
// the data sent by flow or another client will be placed in chunkData.
type Request struct {
	FlowChunkNumber  int32  `json:"flowChunkNumber"`  // The chunk being sent.
	FlowTotalChunks  int32  `json:"flowTotalChunks"`  // The total number of chunks to send.
	FlowChunkSize    int32  `json:"flowChunkSize"`    // The size of the chunk.
	FlowTotalSize    int64  `json:"flowTotalSize"`    // The size of the file being uploaded.
	FlowIdentifier   string `json:"flowIdentifier"`   // A unique identifier used by Flow. Not guaranteed to be a GUID.
	FlowFileName     string `json:"flowFilename"`     // The file name being uploaded.
	FlowRelativePath string `json:"flowRelativePath"` // When available the relative file path.
	ProjectID        string `json:"projectID"`        // Materials Commons Project ID.
	DirectoryID      string `json:"directoryID"`      // Materials Commons Directory ID.
	FileID           string `json:"fileID"`           // Materials Commons File ID.
	Chunk            []byte `json:"-"`                // The file data.
	ChunkHash        string `json:"chunkHash"`        // The computed MD5 hash for the chunk (optional).
	FileHash         string `json:"fileHash"`         // The computed MD5 hash for the file (optional)
}

func (r *Request) UploadID() string {
	return fmt.Sprintf("%s-%s-%s", r.ProjectID, r.DirectoryID, r.FileID)
}

func (r *Request) Write() error {
	path := r.Path()
	err := r.validateWrite(path)
	switch {
	case err == nil:
		return ioutil.WriteFile(path, r.Chunk, ChunkPerms)
	case err == app.ErrExists:
		return nil
	default:
		return err
	}
}

// validateWrite determines if a particular chunk can be written.
// If the size of the ondisk chunk is smaller than the request
// chunk then that chunk is incomplete and we allow a write to it.
func (r *Request) validateWrite(path string) error {
	if err := os.MkdirAll(path, ChunkPerms); err != nil {
		return err
	}
	finfo, err := os.Stat(path)
	switch {
	case os.IsNotExist(err):
		return nil
	case err != nil:
		return app.ErrInvalid
	case finfo.Size() < int64(r.FlowChunkSize):
		return nil
	case finfo.Size() == int64(r.FlowChunkSize):
		return app.ErrExists
	default:
		return app.ErrInvalid
	}
}

func (r *Request) Path() string {
	return filepath.Join(r.Dir(), fmt.Sprintf("%d", r.FlowChunkNumber))
}

func (r *Request) Dir() string {
	mcdir := app.MCDir.Path()
	uploadPath := filepath.Join(mcdir, "upload", r.ProjectID, r.DirectoryID, r.FileID)
	return uploadPath
}
<|endoftext|>


<commit_before>package client

import (
	"fmt"
	neturl "net/url"
	"reflect"
	"strings"
	"time"

	_ "github.com/lib/pq"

	"github.com/jmoiron/sqlx"
	"github.com/sosedoff/pgweb/pkg/command"
	"github.com/sosedoff/pgweb/pkg/connection"
	"github.com/sosedoff/pgweb/pkg/history"
	"github.com/sosedoff/pgweb/pkg/shared"
	"github.com/sosedoff/pgweb/pkg/statements"
)

type Client struct {
	db               *sqlx.DB
	tunnel           *Tunnel
	serverVersion    string
	lastQueryTime    time.Time
	External         bool
	History          []history.Record `json:"history"`
	ConnectionString string           `json:"connection_string"`
}

// Struct to hold table rows browsing options
type RowsOptions struct {
	Where      string // Custom filter
	Offset     int    // Number of rows to skip
	Limit      int    // Number of rows to fetch
	SortColumn string // Column to sort by
	SortOrder  string // Sort direction (ASC, DESC)
}

func getSchemaAndTable(str string) (string, string) {
	chunks := strings.Split(str, ".")
	if len(chunks) == 1 {
		return "public", chunks[0]
	}
	return chunks[0], chunks[1]
}

func New() (*Client, error) {
	str, err := connection.BuildString(command.Opts)

	if command.Opts.Debug && str != "" {
		fmt.Println("Creating a new client for:", str)
	}

	if err != nil {
		return nil, err
	}

	db, err := sqlx.Open("postgres", str)
	if err != nil {
		return nil, err
	}

	client := Client{
		db:               db,
		ConnectionString: str,
		History:          history.New(),
	}

	client.setServerVersion()
	return &client, nil
}

func NewFromUrl(url string, sshInfo *shared.SSHInfo) (*Client, error) {
	var tunnel *Tunnel

	if sshInfo != nil {
		if command.Opts.DisableSSH {
			return nil, fmt.Errorf("ssh connections are disabled")
		}
		if command.Opts.Debug {
			fmt.Println("Opening SSH tunnel for:", sshInfo)
		}

		tunnel, err := NewTunnel(sshInfo, url)
		if err != nil {
			tunnel.Close()
			return nil, err
		}

		err = tunnel.Configure()
		if err != nil {
			tunnel.Close()
			return nil, err
		}

		go tunnel.Start()

		uri, err := neturl.Parse(url)
		if err != nil {
			tunnel.Close()
			return nil, err
		}

		// Override remote postgres port with local proxy port
		url = strings.Replace(url, uri.Host, fmt.Sprintf("127.0.0.1:%v", tunnel.Port), 1)
	}

	if command.Opts.Debug {
		fmt.Println("Creating a new client for:", url)
	}

	db, err := sqlx.Open("postgres", url)
	if err != nil {
		return nil, err
	}

	client := Client{
		db:               db,
		tunnel:           tunnel,
		ConnectionString: url,
		History:          history.New(),
	}

	client.setServerVersion()
	return &client, nil
}

func (client *Client) setServerVersion() {
	res, err := client.query("SELECT version()")
	if err != nil || len(res.Rows) < 1 {
		return
	}

	version := res.Rows[0][0].(string)
	client.serverVersion = strings.Split(version, " ")[1]
}

func (client *Client) Test() error {
	return client.db.Ping()
}

func (client *Client) Info() (*Result, error) {
	return client.query(statements.Info)
}

func (client *Client) Databases() ([]string, error) {
	return client.fetchRows(statements.Databases)
}

func (client *Client) Schemas() ([]string, error) {
	return client.fetchRows(statements.Schemas)
}

func (client *Client) Objects() (*Result, error) {
	return client.query(statements.Objects)
}

func (client *Client) Table(table string) (*Result, error) {
	schema, table := getSchemaAndTable(table)
	return client.query(statements.TableSchema, schema, table)
}

func (client *Client) MaterializedView(name string) (*Result, error) {
	return client.query(statements.MaterializedView, name)
}

func (client *Client) TableRows(table string, opts RowsOptions) (*Result, error) {
	schema, table := getSchemaAndTable(table)
	sql := fmt.Sprintf(`SELECT * FROM "%s"."%s"`, schema, table)

	if opts.Where != "" {
		sql += fmt.Sprintf(" WHERE %s", opts.Where)
	}

	if opts.SortColumn != "" {
		if opts.SortOrder == "" {
			opts.SortOrder = "ASC"
		}

		sql += fmt.Sprintf(` ORDER BY "%s" %s`, opts.SortColumn, opts.SortOrder)
	}

	if opts.Limit > 0 {
		sql += fmt.Sprintf(" LIMIT %d", opts.Limit)
	}

	if opts.Offset > 0 {
		sql += fmt.Sprintf(" OFFSET %d", opts.Offset)
	}

	return client.query(sql)
}

func (client *Client) TableRowsCount(table string, opts RowsOptions) (*Result, error) {
	schema, table := getSchemaAndTable(table)
	sql := fmt.Sprintf(`SELECT COUNT(1) FROM "%s"."%s"`, schema, table)

	if opts.Where != "" {
		sql += fmt.Sprintf(" WHERE %s", opts.Where)
	}

	return client.query(sql)
}

func (client *Client) TableInfo(table string) (*Result, error) {
	return client.query(statements.TableInfo, table)
}

func (client *Client) TableIndexes(table string) (*Result, error) {
	schema, table := getSchemaAndTable(table)
	res, err := client.query(statements.TableIndexes, schema, table)

	if err != nil {
		return nil, err
	}

	return res, err
}

func (client *Client) TableConstraints(table string) (*Result, error) {
	schema, table := getSchemaAndTable(table)
	res, err := client.query(statements.TableConstraints, schema, table)

	if err != nil {
		return nil, err
	}

	return res, err
}

// Returns all active queriers on the server
func (client *Client) Activity() (*Result, error) {
	chunks := strings.Split(client.serverVersion, ".")
	version := strings.Join(chunks[0:2], ".")

	query := statements.Activity[version]
	if query == "" {
		query = statements.Activity["default"]
	}

	return client.query(query)
}

func (client *Client) Query(query string) (*Result, error) {
	res, err := client.query(query)

	// Save history records only if query did not fail
	if err == nil && !client.hasHistoryRecord(query) {
		client.History = append(client.History, history.NewRecord(query))
	}

	return res, err
}

func (client *Client) SetReadOnlyMode() error {
	var value string
	if err := client.db.Get(&value, "SHOW default_transaction_read_only;"); err != nil {
		return err
	}

	if value == "off" {
		_, err := client.db.Exec("SET default_transaction_read_only=on;")
		return err
	}

	return nil
}

func (client *Client) ServerVersion() string {
	return client.serverVersion
}

func (client *Client) query(query string, args ...interface{}) (*Result, error) {
	// Update the last usage time
	defer func() {
		client.lastQueryTime = time.Now().UTC()
	}()

	// We're going to force-set transaction mode on every query.
	// This is needed so that default mode could not be changed by user.
	if command.Opts.ReadOnly {
		if err := client.SetReadOnlyMode(); err != nil {
			return nil, err
		}
	}

	action := strings.ToLower(strings.Split(query, " ")[0])
	if action == "update" || action == "delete" {
		res, err := client.db.Exec(query, args...)
		if err != nil {
			return nil, err
		}

		affected, err := res.RowsAffected()
		if err != nil {
			return nil, err
		}

		result := Result{
			Columns: []string{"Rows Affected"},
			Rows: []Row{
				Row{affected},
			},
		}

		return &result, nil
	}

	rows, err := client.db.Queryx(query, args...)
	if err != nil {
		return nil, err
	}

	defer rows.Close()

	cols, err := rows.Columns()
	if err != nil {
		return nil, err
	}

	// Make sure to never return null colums
	if cols == nil {
		cols = []string{}
	}

	result := Result{
		Columns: cols,
		Rows:    []Row{},
	}

	for rows.Next() {
		obj, err := rows.SliceScan()

		for i, item := range obj {
			if item == nil {
				obj[i] = nil
			} else {
				t := reflect.TypeOf(item).Kind().String()

				if t == "slice" {
					obj[i] = string(item.([]byte))
				}
			}
		}

		if err == nil {
			result.Rows = append(result.Rows, obj)
		}
	}

	result.PrepareBigints()

	return &result, nil
}

// Close database connection
func (client *Client) Close() error {
	if client.tunnel != nil {
		client.tunnel.Close()
	}

	if client.db != nil {
		return client.db.Close()
	}

	return nil
}

func (client *Client) IsIdle() bool {
	if command.Opts.ConnectionIdleTimeoutDisabled {
		return true
	}
	return time.Since(client.lastQueryTime).Minutes() > command.Opts.ConnectionIdleTimeout
}

// Fetch all rows as strings for a single column
func (client *Client) fetchRows(q string) ([]string, error) {
	res, err := client.query(q)

	if err != nil {
		return nil, err
	}

	// Init empty slice so json.Marshal will encode it to "[]" instead of "null"
	results := make([]string, 0)

	for _, row := range res.Rows {
		results = append(results, row[0].(string))
	}

	return results, nil
}

func (client *Client) hasHistoryRecord(query string) bool {
	result := false

	for _, record := range client.History {
		if record.Query == query {
			result = true
			break
		}
	}

	return result
}
<commit_msg>#282 fix disabled state per idle timeout<commit_after>package client

import (
	"fmt"
	neturl "net/url"
	"reflect"
	"strings"
	"time"

	_ "github.com/lib/pq"

	"github.com/jmoiron/sqlx"
	"github.com/sosedoff/pgweb/pkg/command"
	"github.com/sosedoff/pgweb/pkg/connection"
	"github.com/sosedoff/pgweb/pkg/history"
	"github.com/sosedoff/pgweb/pkg/shared"
	"github.com/sosedoff/pgweb/pkg/statements"
)

type Client struct {
	db               *sqlx.DB
	tunnel           *Tunnel
	serverVersion    string
	lastQueryTime    time.Time
	External         bool
	History          []history.Record `json:"history"`
	ConnectionString string           `json:"connection_string"`
}

// Struct to hold table rows browsing options
type RowsOptions struct {
	Where      string // Custom filter
	Offset     int    // Number of rows to skip
	Limit      int    // Number of rows to fetch
	SortColumn string // Column to sort by
	SortOrder  string // Sort direction (ASC, DESC)
}

func getSchemaAndTable(str string) (string, string) {
	chunks := strings.Split(str, ".")
	if len(chunks) == 1 {
		return "public", chunks[0]
	}
	return chunks[0], chunks[1]
}

func New() (*Client, error) {
	str, err := connection.BuildString(command.Opts)

	if command.Opts.Debug && str != "" {
		fmt.Println("Creating a new client for:", str)
	}

	if err != nil {
		return nil, err
	}

	db, err := sqlx.Open("postgres", str)
	if err != nil {
		return nil, err
	}

	client := Client{
		db:               db,
		ConnectionString: str,
		History:          history.New(),
	}

	client.setServerVersion()
	return &client, nil
}

func NewFromUrl(url string, sshInfo *shared.SSHInfo) (*Client, error) {
	var tunnel *Tunnel

	if sshInfo != nil {
		if command.Opts.DisableSSH {
			return nil, fmt.Errorf("ssh connections are disabled")
		}
		if command.Opts.Debug {
			fmt.Println("Opening SSH tunnel for:", sshInfo)
		}

		tunnel, err := NewTunnel(sshInfo, url)
		if err != nil {
			tunnel.Close()
			return nil, err
		}

		err = tunnel.Configure()
		if err != nil {
			tunnel.Close()
			return nil, err
		}

		go tunnel.Start()

		uri, err := neturl.Parse(url)
		if err != nil {
			tunnel.Close()
			return nil, err
		}

		// Override remote postgres port with local proxy port
		url = strings.Replace(url, uri.Host, fmt.Sprintf("127.0.0.1:%v", tunnel.Port), 1)
	}

	if command.Opts.Debug {
		fmt.Println("Creating a new client for:", url)
	}

	db, err := sqlx.Open("postgres", url)
	if err != nil {
		return nil, err
	}

	client := Client{
		db:               db,
		tunnel:           tunnel,
		ConnectionString: url,
		History:          history.New(),
	}

	client.setServerVersion()
	return &client, nil
}

func (client *Client) setServerVersion() {
	res, err := client.query("SELECT version()")
	if err != nil || len(res.Rows) < 1 {
		return
	}

	version := res.Rows[0][0].(string)
	client.serverVersion = strings.Split(version, " ")[1]
}

func (client *Client) Test() error {
	return client.db.Ping()
}

func (client *Client) Info() (*Result, error) {
	return client.query(statements.Info)
}

func (client *Client) Databases() ([]string, error) {
	return client.fetchRows(statements.Databases)
}

func (client *Client) Schemas() ([]string, error) {
	return client.fetchRows(statements.Schemas)
}

func (client *Client) Objects() (*Result, error) {
	return client.query(statements.Objects)
}

func (client *Client) Table(table string) (*Result, error) {
	schema, table := getSchemaAndTable(table)
	return client.query(statements.TableSchema, schema, table)
}

func (client *Client) MaterializedView(name string) (*Result, error) {
	return client.query(statements.MaterializedView, name)
}

func (client *Client) TableRows(table string, opts RowsOptions) (*Result, error) {
	schema, table := getSchemaAndTable(table)
	sql := fmt.Sprintf(`SELECT * FROM "%s"."%s"`, schema, table)

	if opts.Where != "" {
		sql += fmt.Sprintf(" WHERE %s", opts.Where)
	}

	if opts.SortColumn != "" {
		if opts.SortOrder == "" {
			opts.SortOrder = "ASC"
		}

		sql += fmt.Sprintf(` ORDER BY "%s" %s`, opts.SortColumn, opts.SortOrder)
	}

	if opts.Limit > 0 {
		sql += fmt.Sprintf(" LIMIT %d", opts.Limit)
	}

	if opts.Offset > 0 {
		sql += fmt.Sprintf(" OFFSET %d", opts.Offset)
	}

	return client.query(sql)
}

func (client *Client) TableRowsCount(table string, opts RowsOptions) (*Result, error) {
	schema, table := getSchemaAndTable(table)
	sql := fmt.Sprintf(`SELECT COUNT(1) FROM "%s"."%s"`, schema, table)

	if opts.Where != "" {
		sql += fmt.Sprintf(" WHERE %s", opts.Where)
	}

	return client.query(sql)
}

func (client *Client) TableInfo(table string) (*Result, error) {
	return client.query(statements.TableInfo, table)
}

func (client *Client) TableIndexes(table string) (*Result, error) {
	schema, table := getSchemaAndTable(table)
	res, err := client.query(statements.TableIndexes, schema, table)

	if err != nil {
		return nil, err
	}

	return res, err
}

func (client *Client) TableConstraints(table string) (*Result, error) {
	schema, table := getSchemaAndTable(table)
	res, err := client.query(statements.TableConstraints, schema, table)

	if err != nil {
		return nil, err
	}

	return res, err
}

// Returns all active queriers on the server
func (client *Client) Activity() (*Result, error) {
	chunks := strings.Split(client.serverVersion, ".")
	version := strings.Join(chunks[0:2], ".")

	query := statements.Activity[version]
	if query == "" {
		query = statements.Activity["default"]
	}

	return client.query(query)
}

func (client *Client) Query(query string) (*Result, error) {
	res, err := client.query(query)

	// Save history records only if query did not fail
	if err == nil && !client.hasHistoryRecord(query) {
		client.History = append(client.History, history.NewRecord(query))
	}

	return res, err
}

func (client *Client) SetReadOnlyMode() error {
	var value string
	if err := client.db.Get(&value, "SHOW default_transaction_read_only;"); err != nil {
		return err
	}

	if value == "off" {
		_, err := client.db.Exec("SET default_transaction_read_only=on;")
		return err
	}

	return nil
}

func (client *Client) ServerVersion() string {
	return client.serverVersion
}

func (client *Client) query(query string, args ...interface{}) (*Result, error) {
	// Update the last usage time
	defer func() {
		client.lastQueryTime = time.Now().UTC()
	}()

	// We're going to force-set transaction mode on every query.
	// This is needed so that default mode could not be changed by user.
	if command.Opts.ReadOnly {
		if err := client.SetReadOnlyMode(); err != nil {
			return nil, err
		}
	}

	action := strings.ToLower(strings.Split(query, " ")[0])
	if action == "update" || action == "delete" {
		res, err := client.db.Exec(query, args...)
		if err != nil {
			return nil, err
		}

		affected, err := res.RowsAffected()
		if err != nil {
			return nil, err
		}

		result := Result{
			Columns: []string{"Rows Affected"},
			Rows: []Row{
				Row{affected},
			},
		}

		return &result, nil
	}

	rows, err := client.db.Queryx(query, args...)
	if err != nil {
		return nil, err
	}

	defer rows.Close()

	cols, err := rows.Columns()
	if err != nil {
		return nil, err
	}

	// Make sure to never return null colums
	if cols == nil {
		cols = []string{}
	}

	result := Result{
		Columns: cols,
		Rows:    []Row{},
	}

	for rows.Next() {
		obj, err := rows.SliceScan()

		for i, item := range obj {
			if item == nil {
				obj[i] = nil
			} else {
				t := reflect.TypeOf(item).Kind().String()

				if t == "slice" {
					obj[i] = string(item.([]byte))
				}
			}
		}

		if err == nil {
			result.Rows = append(result.Rows, obj)
		}
	}

	result.PrepareBigints()

	return &result, nil
}

// Close database connection
func (client *Client) Close() error {
	if client.tunnel != nil {
		client.tunnel.Close()
	}

	if client.db != nil {
		return client.db.Close()
	}

	return nil
}

func (client *Client) IsIdle() bool {
	if command.Opts.ConnectionIdleTimeoutDisabled {
		return false
	}
	return time.Since(client.lastQueryTime).Minutes() > command.Opts.ConnectionIdleTimeout
}

// Fetch all rows as strings for a single column
func (client *Client) fetchRows(q string) ([]string, error) {
	res, err := client.query(q)

	if err != nil {
		return nil, err
	}

	// Init empty slice so json.Marshal will encode it to "[]" instead of "null"
	results := make([]string, 0)

	for _, row := range res.Rows {
		results = append(results, row[0].(string))
	}

	return results, nil
}

func (client *Client) hasHistoryRecord(query string) bool {
	result := false

	for _, record := range client.History {
		if record.Query == query {
			result = true
			break
		}
	}

	return result
}
<|endoftext|>


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

package client

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"text/tabwriter"
	"time"

	clientapi "github.com/cilium/cilium/api/v1/client"
	"github.com/cilium/cilium/api/v1/models"
	"github.com/cilium/cilium/daemon/defaults"
	"github.com/cilium/cilium/pkg/ip"

	runtime_client "github.com/go-openapi/runtime/client"
	"github.com/go-openapi/strfmt"
)

type Client struct {
	clientapi.Cilium
}

func configureTransport(tr *http.Transport, proto, addr string) *http.Transport {
	if tr == nil {
		tr = &http.Transport{}
	}

	if proto == "unix" {
		// No need for compression in local communications.
		tr.DisableCompression = true
		tr.Dial = func(_, _ string) (net.Conn, error) {
			return net.Dial(proto, addr)
		}
	} else {
		tr.Proxy = http.ProxyFromEnvironment
		tr.Dial = (&net.Dialer{}).Dial
	}

	return tr
}

// NewDefaultClient creates a client with default parameters connecting to UNIX domain socket.
func NewDefaultClient() (*Client, error) {
	return NewClient("")
}

// NewClient creates a client for the given `host`.
func NewClient(host string) (*Client, error) {
	if host == "" {
		// Check if environment variable points to socket
		e := os.Getenv(defaults.SockPathEnv)
		if e == "" {
			// If unset, fall back to default value
			e = defaults.SockPath
		}
		host = "unix://" + e
	}
	tmp := strings.SplitN(host, "://", 2)
	if len(tmp) != 2 {
		return nil, fmt.Errorf("invalid host format '%s'", host)
	}

	switch tmp[0] {
	case "tcp":
		if _, err := url.Parse("tcp://" + tmp[1]); err != nil {
			return nil, err
		}
		host = "http://" + tmp[1]
	case "unix":
		host = tmp[1]
	}

	transport := configureTransport(nil, tmp[0], host)
	httpClient := &http.Client{Transport: transport}
	clientTrans := runtime_client.NewWithClient(tmp[1], clientapi.DefaultBasePath,
		clientapi.DefaultSchemes, httpClient)
	return &Client{*clientapi.New(clientTrans, strfmt.Default)}, nil
}

// Hint tries to improve the error message displayed to the user.
func Hint(err error) error {
	if err == nil {
		return err
	}
	e, _ := url.PathUnescape(err.Error())
	if strings.Contains(err.Error(), defaults.SockPath) {
		return fmt.Errorf("%s\nIs the agent running?", e)
	}
	return fmt.Errorf("%s", e)
}

func formatNodeAddress(w io.Writer, elem *models.NodeAddressingElement, title, prefix string) bool {
	if elem.Enabled || title == "" {
		if title != "" {
			fmt.Fprintf(w, "%s%s Address:\t%s\n", prefix, title, elem.IP)
		} else {
			fmt.Fprintf(w, "%s%s:\n", prefix, elem.IP)
		}
		if elem.AddressType != "" {
			fmt.Fprintf(w, "%s Type:\t%s\n", prefix, elem.AddressType)
		}
		if elem.AllocRange != "" {
			fmt.Fprintf(w, "%sAllocRange:\t%s\n", prefix, elem.AllocRange)
		}

		return true
	}

	return false
}

func timeSince(since time.Time) string {
	out := "never"
	if !since.IsZero() {
		// Poor man's implementtion of time.Truncate(). Can be refined
		// when we rebase to go 1.9
		t := time.Since(since)
		t -= t % time.Second
		out = t.String() + " ago"
	}

	return out
}

func stateUnhealthy(state string) bool {
	return state == models.StatusStateWarning ||
		state == models.StatusStateFailure
}

func statusUnhealthy(s *models.Status) bool {
	if s != nil {
		return stateUnhealthy(s.State)
	}
	return false
}

// FormatStatusResponseBrief writes a one-line status to the writer. If
// everything ok, this is "ok", otherwise a message of the form "error in ..."
func FormatStatusResponseBrief(w io.Writer, sr *models.StatusResponse) {
	msg := ""

	switch {
	case statusUnhealthy(sr.Cilium):
		msg = fmt.Sprintf("cilium: %s", sr.Cilium.Msg)
	case statusUnhealthy(sr.ContainerRuntime):
		msg = fmt.Sprintf("container runtime: %s", sr.ContainerRuntime.Msg)
	case statusUnhealthy(sr.Kvstore):
		msg = fmt.Sprintf("kvstore: %s", sr.Kvstore.Msg)
	case sr.Kubernetes != nil && stateUnhealthy(sr.Kubernetes.State):
		msg = fmt.Sprintf("kubernetes: %s", sr.Kubernetes.Msg)
	case sr.Cluster != nil && statusUnhealthy(sr.Cluster.CiliumHealth):
		msg = fmt.Sprintf("cilium-health: %s", sr.Cluster.CiliumHealth.Msg)
	}

	// Only bother looking at controller failures if everything else is ok
	if msg == "" {
		for _, ctrl := range sr.Controllers {
			if ctrl.Status == nil {
				continue
			}
			if ctrl.Status.LastFailureMsg != "" {
				msg = fmt.Sprintf("controller %s: %s",
					ctrl.Name, ctrl.Status.LastFailureMsg)
				break
			}
		}
	}

	if msg == "" {
		fmt.Fprintf(w, "OK\n")
	} else {
		fmt.Fprintf(w, "error in %s\n", msg)
	}
}

// FormatStatusResponse writes a StatusResponse as a string to the writer.
//
// The parameters 'allAddresses', 'allControllers', 'allNodes', respectively,
// cause all details about that aspect of the status to be printed to the
// terminal. For each of these, if they are false then only a summary will be
// printed, with perhaps some detail if there are errors.
func FormatStatusResponse(w io.Writer, sr *models.StatusResponse, allAddresses, allControllers, allNodes bool) {
	if sr.Kvstore != nil {
		fmt.Fprintf(w, "KVStore:\t%s\t%s\n", sr.Kvstore.State, sr.Kvstore.Msg)
	}
	if sr.ContainerRuntime != nil {
		fmt.Fprintf(w, "ContainerRuntime:\t%s\t%s\n",
			sr.ContainerRuntime.State, sr.ContainerRuntime.Msg)
	}
	if sr.Kubernetes != nil {
		fmt.Fprintf(w, "Kubernetes:\t%s\t%s\n", sr.Kubernetes.State, sr.Kubernetes.Msg)
		if sr.Kubernetes.State != models.K8sStatusStateDisabled {
			fmt.Fprintf(w, "Kubernetes APIs:\t[\"%s\"]\n", strings.Join(sr.Kubernetes.K8sAPIVersions, "\", \""))
		}
	}
	if sr.Cilium != nil {
		fmt.Fprintf(w, "Cilium:\t%s\t%s\n", sr.Cilium.State, sr.Cilium.Msg)
	}

	if nm := sr.NodeMonitor; nm != nil {
		fmt.Fprintf(w, "NodeMonitor:\tListening for events on %d CPUs with %dx%d of shared memory\n",
			nm.Cpus, nm.Npages, nm.Pagesize)
		if nm.Lost != 0 || nm.Unknown != 0 {
			fmt.Fprintf(w, "\t%d events lost, %d unknown notifications\n", nm.Lost, nm.Unknown)
		}
	} else {
		fmt.Fprintf(w, "NodeMonitor:\tDisabled\n")
	}

	var localNode *models.NodeElement
	if sr.Cluster != nil {
		if sr.Cluster.CiliumHealth != nil {
			ch := sr.Cluster.CiliumHealth
			fmt.Fprintf(w, "Cilium health daemon:\t%s\t%s\n", ch.State, ch.Msg)
		}
		for _, node := range sr.Cluster.Nodes {
			if node.Name == sr.Cluster.Self {
				localNode = node
			} else {
				continue
			}
		}
	}

	if sr.IPAM != nil {
		var v4CIDR, v6CIDR string
		if localNode != nil {
			if nIPs := ip.CountIPsInCIDR(localNode.PrimaryAddress.IPV4.AllocRange); nIPs > 0 {
				v4CIDR = fmt.Sprintf("/%d", nIPs)
			}
			if nIPs := ip.CountIPsInCIDR(localNode.PrimaryAddress.IPV6.AllocRange); nIPs > 0 {
				v6CIDR = fmt.Sprintf("/%d", nIPs)
			}
		}
		fmt.Fprintf(w, "IPv4 address pool:\t%d%s allocated\n", len(sr.IPAM.IPV4), v4CIDR)
		if allAddresses {
			for _, ipv4 := range sr.IPAM.IPV4 {
				fmt.Fprintf(w, "  %s\n", ipv4)
			}
		}
		fmt.Fprintf(w, "IPv6 address pool:\t%d%s allocated\n", len(sr.IPAM.IPV6), v6CIDR)
		if allAddresses {
			for _, ipv6 := range sr.IPAM.IPV6 {
				fmt.Fprintf(w, "  %s\n", ipv6)
			}
		}
	}

	if sr.Cluster != nil {
		fmt.Fprintf(w, "Known cluster nodes:\t%d\n", len(sr.Cluster.Nodes))
		for _, node := range sr.Cluster.Nodes {
			localStr := ""
			if node.Name == sr.Cluster.Self {
				localStr = " (localhost)"
			} else if !allNodes {
				continue
			}
			fmt.Fprintf(w, " %s%s:\n", node.Name, localStr)
			formatNodeAddress(w, node.PrimaryAddress.IPV4, "Primary", "  ")
			formatNodeAddress(w, node.PrimaryAddress.IPV6, "Primary", "  ")

			buf := new(bytes.Buffer)
			secondary := false
			fmt.Fprintf(buf, "  Secondary Addresses:\n")
			for _, elem := range node.SecondaryAddresses {
				if formatNodeAddress(buf, elem, "", "   ") {
					secondary = true
				}
			}
			if secondary {
				fmt.Fprintf(w, "%s", buf.String())
			}
			if node.HealthEndpointAddress != nil {
				formatNodeAddress(w, node.HealthEndpointAddress.IPV4, "Health Endpoint", "  ")
				formatNodeAddress(w, node.HealthEndpointAddress.IPV6, "Health Endpoint", "  ")
			}
		}
	}

	if sr.Controllers != nil {
		nFailing, out := 0, []string{"  Name\tLast success\tLast error\tCount\tMessage\n"}
		for _, ctrl := range sr.Controllers {
			status := ctrl.Status
			if status == nil {
				continue
			}

			if status.ConsecutiveFailureCount > 0 {
				nFailing++
			} else if !allControllers {
				continue
			}

			failSince := timeSince(time.Time(status.LastFailureTimestamp))
			successSince := timeSince(time.Time(status.LastSuccessTimestamp))

			err := "no error"
			if status.LastFailureMsg != "" {
				err = status.LastFailureMsg
			}

			out = append(out, fmt.Sprintf("  %s\t%s\t%s\t%d\t%s\t\n",
				ctrl.Name, successSince, failSince, status.ConsecutiveFailureCount, err))
		}

		nOK := len(sr.Controllers) - nFailing
		fmt.Fprintf(w, "Controller Status:\t%d/%d healthy\n", nOK, len(sr.Controllers))
		if len(out) > 1 {
			tab := tabwriter.NewWriter(w, 0, 0, 3, ' ', 0)
			for _, s := range out {
				fmt.Fprint(tab, s)
			}
			tab.Flush()
		}

	}
}
<commit_msg>cilium: Don't print list of nodes in status<commit_after>// Copyright 2016-2017 Authors of Cilium
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package client

import (
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"text/tabwriter"
	"time"

	clientapi "github.com/cilium/cilium/api/v1/client"
	"github.com/cilium/cilium/api/v1/models"
	"github.com/cilium/cilium/daemon/defaults"
	"github.com/cilium/cilium/pkg/ip"

	runtime_client "github.com/go-openapi/runtime/client"
	"github.com/go-openapi/strfmt"
)

type Client struct {
	clientapi.Cilium
}

func configureTransport(tr *http.Transport, proto, addr string) *http.Transport {
	if tr == nil {
		tr = &http.Transport{}
	}

	if proto == "unix" {
		// No need for compression in local communications.
		tr.DisableCompression = true
		tr.Dial = func(_, _ string) (net.Conn, error) {
			return net.Dial(proto, addr)
		}
	} else {
		tr.Proxy = http.ProxyFromEnvironment
		tr.Dial = (&net.Dialer{}).Dial
	}

	return tr
}

// NewDefaultClient creates a client with default parameters connecting to UNIX domain socket.
func NewDefaultClient() (*Client, error) {
	return NewClient("")
}

// NewClient creates a client for the given `host`.
func NewClient(host string) (*Client, error) {
	if host == "" {
		// Check if environment variable points to socket
		e := os.Getenv(defaults.SockPathEnv)
		if e == "" {
			// If unset, fall back to default value
			e = defaults.SockPath
		}
		host = "unix://" + e
	}
	tmp := strings.SplitN(host, "://", 2)
	if len(tmp) != 2 {
		return nil, fmt.Errorf("invalid host format '%s'", host)
	}

	switch tmp[0] {
	case "tcp":
		if _, err := url.Parse("tcp://" + tmp[1]); err != nil {
			return nil, err
		}
		host = "http://" + tmp[1]
	case "unix":
		host = tmp[1]
	}

	transport := configureTransport(nil, tmp[0], host)
	httpClient := &http.Client{Transport: transport}
	clientTrans := runtime_client.NewWithClient(tmp[1], clientapi.DefaultBasePath,
		clientapi.DefaultSchemes, httpClient)
	return &Client{*clientapi.New(clientTrans, strfmt.Default)}, nil
}

// Hint tries to improve the error message displayed to the user.
func Hint(err error) error {
	if err == nil {
		return err
	}
	e, _ := url.PathUnescape(err.Error())
	if strings.Contains(err.Error(), defaults.SockPath) {
		return fmt.Errorf("%s\nIs the agent running?", e)
	}
	return fmt.Errorf("%s", e)
}

func timeSince(since time.Time) string {
	out := "never"
	if !since.IsZero() {
		// Poor man's implementtion of time.Truncate(). Can be refined
		// when we rebase to go 1.9
		t := time.Since(since)
		t -= t % time.Second
		out = t.String() + " ago"
	}

	return out
}

func stateUnhealthy(state string) bool {
	return state == models.StatusStateWarning ||
		state == models.StatusStateFailure
}

func statusUnhealthy(s *models.Status) bool {
	if s != nil {
		return stateUnhealthy(s.State)
	}
	return false
}

// FormatStatusResponseBrief writes a one-line status to the writer. If
// everything ok, this is "ok", otherwise a message of the form "error in ..."
func FormatStatusResponseBrief(w io.Writer, sr *models.StatusResponse) {
	msg := ""

	switch {
	case statusUnhealthy(sr.Cilium):
		msg = fmt.Sprintf("cilium: %s", sr.Cilium.Msg)
	case statusUnhealthy(sr.ContainerRuntime):
		msg = fmt.Sprintf("container runtime: %s", sr.ContainerRuntime.Msg)
	case statusUnhealthy(sr.Kvstore):
		msg = fmt.Sprintf("kvstore: %s", sr.Kvstore.Msg)
	case sr.Kubernetes != nil && stateUnhealthy(sr.Kubernetes.State):
		msg = fmt.Sprintf("kubernetes: %s", sr.Kubernetes.Msg)
	case sr.Cluster != nil && statusUnhealthy(sr.Cluster.CiliumHealth):
		msg = fmt.Sprintf("cilium-health: %s", sr.Cluster.CiliumHealth.Msg)
	}

	// Only bother looking at controller failures if everything else is ok
	if msg == "" {
		for _, ctrl := range sr.Controllers {
			if ctrl.Status == nil {
				continue
			}
			if ctrl.Status.LastFailureMsg != "" {
				msg = fmt.Sprintf("controller %s: %s",
					ctrl.Name, ctrl.Status.LastFailureMsg)
				break
			}
		}
	}

	if msg == "" {
		fmt.Fprintf(w, "OK\n")
	} else {
		fmt.Fprintf(w, "error in %s\n", msg)
	}
}

// FormatStatusResponse writes a StatusResponse as a string to the writer.
//
// The parameters 'allAddresses', 'allControllers', 'allNodes', respectively,
// cause all details about that aspect of the status to be printed to the
// terminal. For each of these, if they are false then only a summary will be
// printed, with perhaps some detail if there are errors.
func FormatStatusResponse(w io.Writer, sr *models.StatusResponse, allAddresses, allControllers, allNodes bool) {
	if sr.Kvstore != nil {
		fmt.Fprintf(w, "KVStore:\t%s\t%s\n", sr.Kvstore.State, sr.Kvstore.Msg)
	}
	if sr.ContainerRuntime != nil {
		fmt.Fprintf(w, "ContainerRuntime:\t%s\t%s\n",
			sr.ContainerRuntime.State, sr.ContainerRuntime.Msg)
	}
	if sr.Kubernetes != nil {
		fmt.Fprintf(w, "Kubernetes:\t%s\t%s\n", sr.Kubernetes.State, sr.Kubernetes.Msg)
		if sr.Kubernetes.State != models.K8sStatusStateDisabled {
			fmt.Fprintf(w, "Kubernetes APIs:\t[\"%s\"]\n", strings.Join(sr.Kubernetes.K8sAPIVersions, "\", \""))
		}
	}
	if sr.Cilium != nil {
		fmt.Fprintf(w, "Cilium:\t%s\t%s\n", sr.Cilium.State, sr.Cilium.Msg)
	}

	if nm := sr.NodeMonitor; nm != nil {
		fmt.Fprintf(w, "NodeMonitor:\tListening for events on %d CPUs with %dx%d of shared memory\n",
			nm.Cpus, nm.Npages, nm.Pagesize)
		if nm.Lost != 0 || nm.Unknown != 0 {
			fmt.Fprintf(w, "\t%d events lost, %d unknown notifications\n", nm.Lost, nm.Unknown)
		}
	} else {
		fmt.Fprintf(w, "NodeMonitor:\tDisabled\n")
	}

	var localNode *models.NodeElement
	if sr.Cluster != nil {
		if sr.Cluster.CiliumHealth != nil {
			ch := sr.Cluster.CiliumHealth
			fmt.Fprintf(w, "Cilium health daemon:\t%s\t%s\n", ch.State, ch.Msg)
		}
		for _, node := range sr.Cluster.Nodes {
			if node.Name == sr.Cluster.Self {
				localNode = node
			} else {
				continue
			}
		}
	}

	if sr.IPAM != nil {
		var v4CIDR, v6CIDR string
		if localNode != nil {
			if nIPs := ip.CountIPsInCIDR(localNode.PrimaryAddress.IPV4.AllocRange); nIPs > 0 {
				v4CIDR = fmt.Sprintf("/%d", nIPs)
			}
			if nIPs := ip.CountIPsInCIDR(localNode.PrimaryAddress.IPV6.AllocRange); nIPs > 0 {
				v6CIDR = fmt.Sprintf("/%d", nIPs)
			}
		}
		fmt.Fprintf(w, "IPv4 address pool:\t%d%s allocated\n", len(sr.IPAM.IPV4), v4CIDR)
		if allAddresses {
			for _, ipv4 := range sr.IPAM.IPV4 {
				fmt.Fprintf(w, "  %s\n", ipv4)
			}
		}
		fmt.Fprintf(w, "IPv6 address pool:\t%d%s allocated\n", len(sr.IPAM.IPV6), v6CIDR)
		if allAddresses {
			for _, ipv6 := range sr.IPAM.IPV6 {
				fmt.Fprintf(w, "  %s\n", ipv6)
			}
		}
	}

	if sr.Controllers != nil {
		nFailing, out := 0, []string{"  Name\tLast success\tLast error\tCount\tMessage\n"}
		for _, ctrl := range sr.Controllers {
			status := ctrl.Status
			if status == nil {
				continue
			}

			if status.ConsecutiveFailureCount > 0 {
				nFailing++
			} else if !allControllers {
				continue
			}

			failSince := timeSince(time.Time(status.LastFailureTimestamp))
			successSince := timeSince(time.Time(status.LastSuccessTimestamp))

			err := "no error"
			if status.LastFailureMsg != "" {
				err = status.LastFailureMsg
			}

			out = append(out, fmt.Sprintf("  %s\t%s\t%s\t%d\t%s\t\n",
				ctrl.Name, successSince, failSince, status.ConsecutiveFailureCount, err))
		}

		nOK := len(sr.Controllers) - nFailing
		fmt.Fprintf(w, "Controller Status:\t%d/%d healthy\n", nOK, len(sr.Controllers))
		if len(out) > 1 {
			tab := tabwriter.NewWriter(w, 0, 0, 3, ' ', 0)
			for _, s := range out {
				fmt.Fprint(tab, s)
			}
			tab.Flush()
		}

	}
}
<|endoftext|>