<commit_before>/*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 installer

import (
	"github.com/skyrings/bigfin/backend"
	"github.com/skyrings/skyring-common/provisioner"
	"github.com/skyrings/skyring-common/tools/logger"
	"github.com/skyrings/skyring-common/tools/uuid"

	bigfin_conf "github.com/skyrings/bigfin/conf"
)

type Installer struct {
}

var installer provisioner.Provisioner

func New() (backend.Backend, error) {
	installerapi := new(Installer)
	//Initilaize the provisioner
	installer, err := provisioner.InitializeProvisioner(bigfin_conf.SystemConfig.Provisioners[bigfin_conf.ProviderName])
	if err != nil {
		logger.Get().Error("Unable to initialize the provisioner, skipping the provider:%v", bigfin_conf.SystemConfig.Provisioners[bigfin_conf.ProviderName])
		return installerapi, err
	}
	logger.Get().Debug("Installer:%v", installer)
	return installerapi, nil
}

func (c Installer) CreateCluster(clusterName string, fsid uuid.UUID, mons []backend.Mon, ctxt string) (bool, error) {
	return true, nil
}

func (c Installer) AddMon(clusterName string, mons []backend.Mon, ctxt string) (bool, error) {
	return true, nil
}

func (c Installer) StartMon(nodes []string, ctxt string) (bool, error) {
	return true, nil
}

func (c Installer) AddOSD(clusterName string, osd backend.OSD, ctxt string) (map[string][]string, error) {
	return map[string][]string{}, nil
}

func (c Installer) CreatePool(name string, mon string, clusterName string, pgnum uint, replicas int, quotaMaxObjects int, quotaMaxBytes uint64) (bool, error) {
	return true, nil
}

func (c Installer) ListPoolNames(mon string, clusterName string) ([]string, error) {
	return []string{}, nil
}

func (c Installer) GetClusterStatus(mon string, clusterName string) (status string, err error) {
	return "", nil
}

func (c Installer) GetClusterStats(mon string, clusterName string) (map[string]int64, error) {
	return nil, nil
}

func (c Installer) GetPools(mon string, clusterId uuid.UUID) ([]backend.CephPool, error) {
	return []backend.CephPool{}, nil
}

func (c Installer) UpdatePool(mon string, clusterId uuid.UUID, poolId int, pool map[string]interface{}) (bool, error) {
	return true, nil
}

func (c Installer) RemovePool(mon string, clusterId uuid.UUID, clusterName string, pool string, poolId int, ctxt string) (bool, error) {
	return true, nil
}

func (c Installer) GetOSDDetails(mon string, clusterName string) (osds []backend.OSDDetails, err error) {
	return []backend.OSDDetails{}, nil
}

func (c Installer) GetObjectCount(mon string, clusterName string) (string, error) {
	return "", nil
}

func (c Installer) GetPGSummary(mon string, clusterId uuid.UUID) (backend.PgSummary, error) {
	return backend.PgSummary{}, nil
}

func (c Installer) ExecCmd(mon string, clusterId uuid.UUID, cmd string) (bool, error) {
	return true, nil
}

func (c Installer) GetOSDs(mon string, clusterId uuid.UUID) ([]backend.CephOSD, error) {
	return []backend.CephOSD{}, nil
}

func (c Installer) UpdateOSD(mon string, clusterId uuid.UUID, osdId string, params map[string]interface{}) (bool, error) {
	return true, nil
}

func (c Installer) GetOSD(mon string, clusterId uuid.UUID, osdId string) (backend.CephOSD, error) {
	return backend.CephOSD{}, nil
}
<commit_msg>bigfin: Corrected compilation error<commit_after>/*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 installer

import (
	"github.com/skyrings/bigfin/backend"
	"github.com/skyrings/skyring-common/provisioner"
	"github.com/skyrings/skyring-common/tools/logger"
	"github.com/skyrings/skyring-common/tools/uuid"

	bigfin_conf "github.com/skyrings/bigfin/conf"
)

type Installer struct {
}

var installer provisioner.Provisioner

func New() (backend.Backend, error) {
	installerapi := new(Installer)
	//Initilaize the provisioner
	installer, err := provisioner.InitializeProvisioner(bigfin_conf.SystemConfig.Provisioners[bigfin_conf.ProviderName])
	if err != nil {
		logger.Get().Error("Unable to initialize the provisioner, skipping the provider:%v", bigfin_conf.SystemConfig.Provisioners[bigfin_conf.ProviderName])
		return installerapi, err
	}
	logger.Get().Debug("Installer:%v", installer)
	return installerapi, nil
}

func (c Installer) CreateCluster(clusterName string, fsid uuid.UUID, mons []backend.Mon, ctxt string) (bool, error) {
	return true, nil
}

func (c Installer) AddMon(clusterName string, mons []backend.Mon, ctxt string) (bool, error) {
	return true, nil
}

func (c Installer) StartMon(nodes []string, ctxt string) (bool, error) {
	return true, nil
}

func (c Installer) AddOSD(clusterName string, osd backend.OSD, ctxt string) (map[string][]string, error) {
	return map[string][]string{}, nil
}

func (c Installer) CreatePool(name string, mon string, clusterName string, pgnum uint, replicas int, quotaMaxObjects int, quotaMaxBytes uint64) (bool, error) {
	return true, nil
}

func (c Installer) ListPoolNames(mon string, clusterName string) ([]string, error) {
	return []string{}, nil
}

func (c Installer) GetClusterStatus(mon string, clusterName string) (status string, err error) {
	return "", nil
}

func (c Installer) GetClusterStats(mon string, clusterName string) (map[string]int64, error) {
	return nil, nil
}

func (c Installer) GetPools(mon string, clusterId uuid.UUID) ([]backend.CephPool, error) {
	return []backend.CephPool{}, nil
}

func (c Installer) UpdatePool(mon string, clusterId uuid.UUID, poolId int, pool map[string]interface{}) (bool, error) {
	return true, nil
}

func (c Installer) RemovePool(mon string, clusterId uuid.UUID, clusterName string, pool string, poolId int, ctxt string) (bool, error) {
	return true, nil
}

func (c Installer) GetOSDDetails(mon string, clusterName string) (osds []backend.OSDDetails, err error) {
	return []backend.OSDDetails{}, nil
}

func (c Installer) GetObjectCount(mon string, clusterName string) (string, error) {
	return "", nil
}

func (c Installer) GetPGSummary(mon string, clusterId uuid.UUID) (backend.PgSummary, error) {
	return backend.PgSummary{}, nil
}

func (c Installer) ExecCmd(mon string, clusterId uuid.UUID, cmd string) (bool, string, error) {
	return true, "", nil
}

func (c Installer) GetOSDs(mon string, clusterId uuid.UUID) ([]backend.CephOSD, error) {
	return []backend.CephOSD{}, nil
}

func (c Installer) UpdateOSD(mon string, clusterId uuid.UUID, osdId string, params map[string]interface{}) (bool, error) {
	return true, nil
}

func (c Installer) GetOSD(mon string, clusterId uuid.UUID, osdId string) (backend.CephOSD, error) {
	return backend.CephOSD{}, nil
}
<|endoftext|>


<commit_before>package nomad

import (
	"errors"
	"fmt"
	"io"
	"net"
	"strings"
	"time"

	metrics "github.com/armon/go-metrics"
	"github.com/hashicorp/nomad/acl"
	cstructs "github.com/hashicorp/nomad/client/structs"
	"github.com/hashicorp/nomad/helper"
	"github.com/hashicorp/nomad/nomad/state"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/ugorji/go/codec"
)

// FileSystem endpoint is used for accessing the logs and filesystem of
// allocations from a Node.
type FileSystem struct {
	srv *Server
}

func (f *FileSystem) register() {
	f.srv.streamingRpcs.Register("FileSystem.Logs", f.logs)
	f.srv.streamingRpcs.Register("FileSystem.Stream", f.stream)
}

// handleStreamResultError is a helper for sending an error with a potential
// error code. The transmission of the error is ignored if the error has been
// generated by the closing of the underlying transport.
func (f *FileSystem) handleStreamResultError(err error, code *int64, encoder *codec.Encoder) {
	// Nothing to do as the conn is closed
	if err == io.EOF || strings.Contains(err.Error(), "closed") {
		return
	}

	// Attempt to send the error
	encoder.Encode(&cstructs.StreamErrWrapper{
		Error: cstructs.NewRpcError(err, code),
	})
}

// findNodeConnAndForward is a helper for finding the server with a connection
// to the given node and forwarding the RPC to the correct server. This does not
// work for streaming RPCs.
func (f *FileSystem) findNodeConnAndForward(snap *state.StateSnapshot,
	nodeID, method string, args, reply interface{}) error {

	node, err := snap.NodeByID(nil, nodeID)
	if err != nil {
		return err
	}

	if node == nil {
		return fmt.Errorf("Unknown node %q", nodeID)
	}

	// Determine the Server that has a connection to the node.
	srv, err := f.srv.serverWithNodeConn(nodeID, f.srv.Region())
	if err != nil {
		return err
	}

	if srv == nil {
		return structs.ErrNoNodeConn
	}

	return f.srv.forwardServer(srv, method, args, reply)
}

// List is used to list the contents of an allocation's directory.
func (f *FileSystem) List(args *cstructs.FsListRequest, reply *cstructs.FsListResponse) error {
	// We only allow stale reads since the only potentially stale information is
	// the Node registration and the cost is fairly high for adding another hope
	// in the forwarding chain.
	args.QueryOptions.AllowStale = true

	// Potentially forward to a different region.
	if done, err := f.srv.forward("FileSystem.List", args, args, reply); done {
		return err
	}
	defer metrics.MeasureSince([]string{"nomad", "file_system", "list"}, time.Now())

	// Check filesystem read permissions
	if aclObj, err := f.srv.ResolveToken(args.AuthToken); err != nil {
		return err
	} else if aclObj != nil && !aclObj.AllowNsOp(args.Namespace, acl.NamespaceCapabilityReadFS) {
		return structs.ErrPermissionDenied
	}

	// Verify the arguments.
	if args.AllocID == "" {
		return errors.New("missing allocation ID")
	}

	// Lookup the allocation
	snap, err := f.srv.State().Snapshot()
	if err != nil {
		return err
	}

	alloc, err := snap.AllocByID(nil, args.AllocID)
	if err != nil {
		return err
	}
	if alloc == nil {
		return fmt.Errorf("unknown allocation %q", args.AllocID)
	}

	// Get the connection to the client
	state, ok := f.srv.getNodeConn(alloc.NodeID)
	if !ok {
		return f.findNodeConnAndForward(snap, alloc.NodeID, "FileSystem.List", args, reply)
	}

	// Make the RPC
	return NodeRpc(state.Session, "FileSystem.List", args, reply)
}

// Stat is used to stat a file in the allocation's directory.
func (f *FileSystem) Stat(args *cstructs.FsStatRequest, reply *cstructs.FsStatResponse) error {
	// We only allow stale reads since the only potentially stale information is
	// the Node registration and the cost is fairly high for adding another hope
	// in the forwarding chain.
	args.QueryOptions.AllowStale = true

	// Potentially forward to a different region.
	if done, err := f.srv.forward("FileSystem.Stat", args, args, reply); done {
		return err
	}
	defer metrics.MeasureSince([]string{"nomad", "file_system", "stat"}, time.Now())

	// Check filesystem read permissions
	if aclObj, err := f.srv.ResolveToken(args.AuthToken); err != nil {
		return err
	} else if aclObj != nil && !aclObj.AllowNsOp(args.Namespace, acl.NamespaceCapabilityReadFS) {
		return structs.ErrPermissionDenied
	}

	// Verify the arguments.
	if args.AllocID == "" {
		return errors.New("missing allocation ID")
	}

	// Lookup the allocation
	snap, err := f.srv.State().Snapshot()
	if err != nil {
		return err
	}

	alloc, err := snap.AllocByID(nil, args.AllocID)
	if err != nil {
		return err
	}
	if alloc == nil {
		return fmt.Errorf("unknown allocation %q", args.AllocID)
	}

	// Get the connection to the client
	state, ok := f.srv.getNodeConn(alloc.NodeID)
	if !ok {
		return f.findNodeConnAndForward(snap, alloc.NodeID, "FileSystem.Stat", args, reply)
	}

	// Make the RPC
	return NodeRpc(state.Session, "FileSystem.Stat", args, reply)
}

// stream is is used to stream the contents of file in an allocation's
// directory.
func (f *FileSystem) stream(conn io.ReadWriteCloser) {
	defer conn.Close()
	defer metrics.MeasureSince([]string{"nomad", "file_system", "stream"}, time.Now())

	// Decode the arguments
	var args cstructs.FsStreamRequest
	decoder := codec.NewDecoder(conn, structs.MsgpackHandle)
	encoder := codec.NewEncoder(conn, structs.MsgpackHandle)

	if err := decoder.Decode(&args); err != nil {
		f.handleStreamResultError(err, helper.Int64ToPtr(500), encoder)
		return
	}

	// Check if we need to forward to a different region
	if r := args.RequestRegion(); r != f.srv.Region() {
		// Request the allocation from the target region
		allocReq := &structs.AllocSpecificRequest{
			AllocID:      args.AllocID,
			QueryOptions: args.QueryOptions,
		}
		var allocResp structs.SingleAllocResponse
		if err := f.srv.forwardRegion(r, "Alloc.GetAlloc", allocReq, &allocResp); err != nil {
			f.handleStreamResultError(err, nil, encoder)
			return
		}

		if allocResp.Alloc == nil {
			f.handleStreamResultError(fmt.Errorf("unknown allocation %q", args.AllocID), nil, encoder)
			return
		}

		// Determine the Server that has a connection to the node.
		srv, err := f.srv.serverWithNodeConn(allocResp.Alloc.NodeID, r)
		if err != nil {
			f.handleStreamResultError(err, nil, encoder)
			return
		}

		// Get a connection to the server
		srvConn, err := f.srv.streamingRpc(srv, "FileSystem.Stream")
		if err != nil {
			f.handleStreamResultError(err, nil, encoder)
			return
		}
		defer srvConn.Close()

		// Send the request.
		outEncoder := codec.NewEncoder(srvConn, structs.MsgpackHandle)
		if err := outEncoder.Encode(args); err != nil {
			f.handleStreamResultError(err, nil, encoder)
			return
		}

		structs.Bridge(conn, srvConn)
		return

	}

	// Check node read permissions
	if aclObj, err := f.srv.ResolveToken(args.AuthToken); err != nil {
		f.handleStreamResultError(err, nil, encoder)
		return
	} else if aclObj != nil && !aclObj.AllowNsOp(args.Namespace, acl.NamespaceCapabilityReadFS) {
		f.handleStreamResultError(structs.ErrPermissionDenied, nil, encoder)
		return
	}

	// Verify the arguments.
	if args.AllocID == "" {
		f.handleStreamResultError(errors.New("missing AllocID"), helper.Int64ToPtr(400), encoder)
		return
	}

	// Retrieve the allocation
	snap, err := f.srv.State().Snapshot()
	if err != nil {
		f.handleStreamResultError(err, nil, encoder)
		return
	}

	alloc, err := snap.AllocByID(nil, args.AllocID)
	if err != nil {
		f.handleStreamResultError(err, nil, encoder)
		return
	}
	if alloc == nil {
		f.handleStreamResultError(fmt.Errorf("unknown alloc ID %q", args.AllocID), helper.Int64ToPtr(404), encoder)
		return
	}
	nodeID := alloc.NodeID

	// Get the connection to the client either by forwarding to another server
	// or creating a direct stream
	var clientConn net.Conn
	state, ok := f.srv.getNodeConn(nodeID)
	if !ok {
		// Determine the Server that has a connection to the node.
		srv, err := f.srv.serverWithNodeConn(nodeID, f.srv.Region())
		if err != nil {
			f.handleStreamResultError(err, nil, encoder)
			return
		}

		// Get a connection to the server
		conn, err := f.srv.streamingRpc(srv, "FileSystem.Stream")
		if err != nil {
			f.handleStreamResultError(err, nil, encoder)
			return
		}

		clientConn = conn
	} else {
		stream, err := NodeStreamingRpc(state.Session, "FileSystem.Stream")
		if err != nil {
			f.handleStreamResultError(err, nil, encoder)
			return
		}
		clientConn = stream
	}
	defer clientConn.Close()

	// Send the request.
	outEncoder := codec.NewEncoder(clientConn, structs.MsgpackHandle)
	if err := outEncoder.Encode(args); err != nil {
		f.handleStreamResultError(err, nil, encoder)
		return
	}

	structs.Bridge(conn, clientConn)
	return
}

// logs is used to access an task's logs for a given allocation
func (f *FileSystem) logs(conn io.ReadWriteCloser) {
	defer conn.Close()
	defer metrics.MeasureSince([]string{"nomad", "file_system", "logs"}, time.Now())

	// Decode the arguments
	var args cstructs.FsLogsRequest
	decoder := codec.NewDecoder(conn, structs.MsgpackHandle)
	encoder := codec.NewEncoder(conn, structs.MsgpackHandle)

	if err := decoder.Decode(&args); err != nil {
		f.handleStreamResultError(err, helper.Int64ToPtr(500), encoder)
		return
	}

	// Check if we need to forward to a different region
	if r := args.RequestRegion(); r != f.srv.Region() {
		// Request the allocation from the target region
		allocReq := &structs.AllocSpecificRequest{
			AllocID:      args.AllocID,
			QueryOptions: args.QueryOptions,
		}
		var allocResp structs.SingleAllocResponse
		if err := f.srv.forwardRegion(r, "Alloc.GetAlloc", allocReq, &allocResp); err != nil {
			f.handleStreamResultError(err, nil, encoder)
			return
		}

		if allocResp.Alloc == nil {
			f.handleStreamResultError(fmt.Errorf("unknown allocation %q", args.AllocID), nil, encoder)
			return
		}

		// Determine the Server that has a connection to the node.
		srv, err := f.srv.serverWithNodeConn(allocResp.Alloc.NodeID, r)
		if err != nil {
			f.handleStreamResultError(err, nil, encoder)
			return
		}

		// Get a connection to the server
		srvConn, err := f.srv.streamingRpc(srv, "FileSystem.Logs")
		if err != nil {
			f.handleStreamResultError(err, nil, encoder)
			return
		}
		defer srvConn.Close()

		// Send the request.
		outEncoder := codec.NewEncoder(srvConn, structs.MsgpackHandle)
		if err := outEncoder.Encode(args); err != nil {
			f.handleStreamResultError(err, nil, encoder)
			return
		}

		structs.Bridge(conn, srvConn)
		return

	}

	// Check node read permissions
	if aclObj, err := f.srv.ResolveToken(args.AuthToken); err != nil {
		f.handleStreamResultError(err, nil, encoder)
		return
	} else if aclObj != nil {
		readfs := aclObj.AllowNsOp(args.QueryOptions.Namespace, acl.NamespaceCapabilityReadFS)
		logs := aclObj.AllowNsOp(args.QueryOptions.Namespace, acl.NamespaceCapabilityReadLogs)
		if !readfs && !logs {
			f.handleStreamResultError(structs.ErrPermissionDenied, nil, encoder)
			return
		}
	}

	// Verify the arguments.
	if args.AllocID == "" {
		f.handleStreamResultError(errors.New("missing AllocID"), helper.Int64ToPtr(400), encoder)
		return
	}

	// Retrieve the allocation
	snap, err := f.srv.State().Snapshot()
	if err != nil {
		f.handleStreamResultError(err, nil, encoder)
		return
	}

	alloc, err := snap.AllocByID(nil, args.AllocID)
	if err != nil {
		f.handleStreamResultError(err, nil, encoder)
		return
	}
	if alloc == nil {
		f.handleStreamResultError(fmt.Errorf("unknown alloc ID %q", args.AllocID), helper.Int64ToPtr(404), encoder)
		return
	}
	nodeID := alloc.NodeID

	// Get the connection to the client either by forwarding to another server
	// or creating a direct stream
	var clientConn net.Conn
	state, ok := f.srv.getNodeConn(nodeID)
	if !ok {
		// Determine the Server that has a connection to the node.
		srv, err := f.srv.serverWithNodeConn(nodeID, f.srv.Region())
		if err != nil {
			f.handleStreamResultError(err, nil, encoder)
			return
		}

		// Get a connection to the server
		conn, err := f.srv.streamingRpc(srv, "FileSystem.Logs")
		if err != nil {
			f.handleStreamResultError(err, nil, encoder)
			return
		}

		clientConn = conn
	} else {
		stream, err := NodeStreamingRpc(state.Session, "FileSystem.Logs")
		if err != nil {
			f.handleStreamResultError(err, nil, encoder)
			return
		}
		clientConn = stream
	}
	defer clientConn.Close()

	// Send the request.
	outEncoder := codec.NewEncoder(clientConn, structs.MsgpackHandle)
	if err := outEncoder.Encode(args); err != nil {
		f.handleStreamResultError(err, nil, encoder)
		return
	}

	structs.Bridge(conn, clientConn)
	return
}
<commit_msg>Streaming helper<commit_after>package nomad

import (
	"errors"
	"fmt"
	"io"
	"net"
	"strings"
	"time"

	metrics "github.com/armon/go-metrics"
	"github.com/hashicorp/nomad/acl"
	cstructs "github.com/hashicorp/nomad/client/structs"
	"github.com/hashicorp/nomad/helper"
	"github.com/hashicorp/nomad/nomad/state"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/ugorji/go/codec"
)

// FileSystem endpoint is used for accessing the logs and filesystem of
// allocations from a Node.
type FileSystem struct {
	srv *Server
}

func (f *FileSystem) register() {
	f.srv.streamingRpcs.Register("FileSystem.Logs", f.logs)
	f.srv.streamingRpcs.Register("FileSystem.Stream", f.stream)
}

// handleStreamResultError is a helper for sending an error with a potential
// error code. The transmission of the error is ignored if the error has been
// generated by the closing of the underlying transport.
func (f *FileSystem) handleStreamResultError(err error, code *int64, encoder *codec.Encoder) {
	// Nothing to do as the conn is closed
	if err == io.EOF || strings.Contains(err.Error(), "closed") {
		return
	}

	// Attempt to send the error
	encoder.Encode(&cstructs.StreamErrWrapper{
		Error: cstructs.NewRpcError(err, code),
	})
}

// findNodeConnAndForward is a helper for finding the server with a connection
// to the given node and forwarding the RPC to the correct server. This does not
// work for streaming RPCs.
func (f *FileSystem) findNodeConnAndForward(snap *state.StateSnapshot,
	nodeID, method string, args, reply interface{}) error {

	node, err := snap.NodeByID(nil, nodeID)
	if err != nil {
		return err
	}

	if node == nil {
		return fmt.Errorf("Unknown node %q", nodeID)
	}

	// Determine the Server that has a connection to the node.
	srv, err := f.srv.serverWithNodeConn(nodeID, f.srv.Region())
	if err != nil {
		return err
	}

	if srv == nil {
		return structs.ErrNoNodeConn
	}

	return f.srv.forwardServer(srv, method, args, reply)
}

// forwardRegionStreamingRpc is used to make a streaming RPC to a different
// region. It looks up the allocation in the remote region to determine what
// remote server can route the request.
func (f *FileSystem) forwardRegionStreamingRpc(conn io.ReadWriteCloser,
	encoder *codec.Encoder, args interface{}, method, allocID string, qo *structs.QueryOptions) {
	// Request the allocation from the target region
	allocReq := &structs.AllocSpecificRequest{
		AllocID:      allocID,
		QueryOptions: *qo,
	}
	var allocResp structs.SingleAllocResponse
	if err := f.srv.forwardRegion(qo.RequestRegion(), "Alloc.GetAlloc", allocReq, &allocResp); err != nil {
		f.handleStreamResultError(err, nil, encoder)
		return
	}

	if allocResp.Alloc == nil {
		f.handleStreamResultError(fmt.Errorf("unknown allocation %q", allocID), nil, encoder)
		return
	}

	// Determine the Server that has a connection to the node.
	srv, err := f.srv.serverWithNodeConn(allocResp.Alloc.NodeID, qo.RequestRegion())
	if err != nil {
		f.handleStreamResultError(err, nil, encoder)
		return
	}

	// Get a connection to the server
	srvConn, err := f.srv.streamingRpc(srv, method)
	if err != nil {
		f.handleStreamResultError(err, nil, encoder)
		return
	}
	defer srvConn.Close()

	// Send the request.
	outEncoder := codec.NewEncoder(srvConn, structs.MsgpackHandle)
	if err := outEncoder.Encode(args); err != nil {
		f.handleStreamResultError(err, nil, encoder)
		return
	}

	structs.Bridge(conn, srvConn)
}

// List is used to list the contents of an allocation's directory.
func (f *FileSystem) List(args *cstructs.FsListRequest, reply *cstructs.FsListResponse) error {
	// We only allow stale reads since the only potentially stale information is
	// the Node registration and the cost is fairly high for adding another hope
	// in the forwarding chain.
	args.QueryOptions.AllowStale = true

	// Potentially forward to a different region.
	if done, err := f.srv.forward("FileSystem.List", args, args, reply); done {
		return err
	}
	defer metrics.MeasureSince([]string{"nomad", "file_system", "list"}, time.Now())

	// Check filesystem read permissions
	if aclObj, err := f.srv.ResolveToken(args.AuthToken); err != nil {
		return err
	} else if aclObj != nil && !aclObj.AllowNsOp(args.Namespace, acl.NamespaceCapabilityReadFS) {
		return structs.ErrPermissionDenied
	}

	// Verify the arguments.
	if args.AllocID == "" {
		return errors.New("missing allocation ID")
	}

	// Lookup the allocation
	snap, err := f.srv.State().Snapshot()
	if err != nil {
		return err
	}

	alloc, err := snap.AllocByID(nil, args.AllocID)
	if err != nil {
		return err
	}
	if alloc == nil {
		return fmt.Errorf("unknown allocation %q", args.AllocID)
	}

	// Get the connection to the client
	state, ok := f.srv.getNodeConn(alloc.NodeID)
	if !ok {
		return f.findNodeConnAndForward(snap, alloc.NodeID, "FileSystem.List", args, reply)
	}

	// Make the RPC
	return NodeRpc(state.Session, "FileSystem.List", args, reply)
}

// Stat is used to stat a file in the allocation's directory.
func (f *FileSystem) Stat(args *cstructs.FsStatRequest, reply *cstructs.FsStatResponse) error {
	// We only allow stale reads since the only potentially stale information is
	// the Node registration and the cost is fairly high for adding another hope
	// in the forwarding chain.
	args.QueryOptions.AllowStale = true

	// Potentially forward to a different region.
	if done, err := f.srv.forward("FileSystem.Stat", args, args, reply); done {
		return err
	}
	defer metrics.MeasureSince([]string{"nomad", "file_system", "stat"}, time.Now())

	// Check filesystem read permissions
	if aclObj, err := f.srv.ResolveToken(args.AuthToken); err != nil {
		return err
	} else if aclObj != nil && !aclObj.AllowNsOp(args.Namespace, acl.NamespaceCapabilityReadFS) {
		return structs.ErrPermissionDenied
	}

	// Verify the arguments.
	if args.AllocID == "" {
		return errors.New("missing allocation ID")
	}

	// Lookup the allocation
	snap, err := f.srv.State().Snapshot()
	if err != nil {
		return err
	}

	alloc, err := snap.AllocByID(nil, args.AllocID)
	if err != nil {
		return err
	}
	if alloc == nil {
		return fmt.Errorf("unknown allocation %q", args.AllocID)
	}

	// Get the connection to the client
	state, ok := f.srv.getNodeConn(alloc.NodeID)
	if !ok {
		return f.findNodeConnAndForward(snap, alloc.NodeID, "FileSystem.Stat", args, reply)
	}

	// Make the RPC
	return NodeRpc(state.Session, "FileSystem.Stat", args, reply)
}

// stream is is used to stream the contents of file in an allocation's
// directory.
func (f *FileSystem) stream(conn io.ReadWriteCloser) {
	defer conn.Close()
	defer metrics.MeasureSince([]string{"nomad", "file_system", "stream"}, time.Now())

	// Decode the arguments
	var args cstructs.FsStreamRequest
	decoder := codec.NewDecoder(conn, structs.MsgpackHandle)
	encoder := codec.NewEncoder(conn, structs.MsgpackHandle)

	if err := decoder.Decode(&args); err != nil {
		f.handleStreamResultError(err, helper.Int64ToPtr(500), encoder)
		return
	}

	// Check if we need to forward to a different region
	if r := args.RequestRegion(); r != f.srv.Region() {
		f.forwardRegionStreamingRpc(conn, encoder, &args, "FileSystem.Stream",
			args.AllocID, &args.QueryOptions)
		return
	}

	// Check node read permissions
	if aclObj, err := f.srv.ResolveToken(args.AuthToken); err != nil {
		f.handleStreamResultError(err, nil, encoder)
		return
	} else if aclObj != nil && !aclObj.AllowNsOp(args.Namespace, acl.NamespaceCapabilityReadFS) {
		f.handleStreamResultError(structs.ErrPermissionDenied, nil, encoder)
		return
	}

	// Verify the arguments.
	if args.AllocID == "" {
		f.handleStreamResultError(errors.New("missing AllocID"), helper.Int64ToPtr(400), encoder)
		return
	}

	// Retrieve the allocation
	snap, err := f.srv.State().Snapshot()
	if err != nil {
		f.handleStreamResultError(err, nil, encoder)
		return
	}

	alloc, err := snap.AllocByID(nil, args.AllocID)
	if err != nil {
		f.handleStreamResultError(err, nil, encoder)
		return
	}
	if alloc == nil {
		f.handleStreamResultError(fmt.Errorf("unknown alloc ID %q", args.AllocID), helper.Int64ToPtr(404), encoder)
		return
	}
	nodeID := alloc.NodeID

	// Get the connection to the client either by forwarding to another server
	// or creating a direct stream
	var clientConn net.Conn
	state, ok := f.srv.getNodeConn(nodeID)
	if !ok {
		// Determine the Server that has a connection to the node.
		srv, err := f.srv.serverWithNodeConn(nodeID, f.srv.Region())
		if err != nil {
			f.handleStreamResultError(err, nil, encoder)
			return
		}

		// Get a connection to the server
		conn, err := f.srv.streamingRpc(srv, "FileSystem.Stream")
		if err != nil {
			f.handleStreamResultError(err, nil, encoder)
			return
		}

		clientConn = conn
	} else {
		stream, err := NodeStreamingRpc(state.Session, "FileSystem.Stream")
		if err != nil {
			f.handleStreamResultError(err, nil, encoder)
			return
		}
		clientConn = stream
	}
	defer clientConn.Close()

	// Send the request.
	outEncoder := codec.NewEncoder(clientConn, structs.MsgpackHandle)
	if err := outEncoder.Encode(args); err != nil {
		f.handleStreamResultError(err, nil, encoder)
		return
	}

	structs.Bridge(conn, clientConn)
	return
}

// logs is used to access an task's logs for a given allocation
func (f *FileSystem) logs(conn io.ReadWriteCloser) {
	defer conn.Close()
	defer metrics.MeasureSince([]string{"nomad", "file_system", "logs"}, time.Now())

	// Decode the arguments
	var args cstructs.FsLogsRequest
	decoder := codec.NewDecoder(conn, structs.MsgpackHandle)
	encoder := codec.NewEncoder(conn, structs.MsgpackHandle)

	if err := decoder.Decode(&args); err != nil {
		f.handleStreamResultError(err, helper.Int64ToPtr(500), encoder)
		return
	}

	// Check if we need to forward to a different region
	if r := args.RequestRegion(); r != f.srv.Region() {
		f.forwardRegionStreamingRpc(conn, encoder, &args, "FileSystem.Logs",
			args.AllocID, &args.QueryOptions)
		return
	}

	// Check node read permissions
	if aclObj, err := f.srv.ResolveToken(args.AuthToken); err != nil {
		f.handleStreamResultError(err, nil, encoder)
		return
	} else if aclObj != nil {
		readfs := aclObj.AllowNsOp(args.QueryOptions.Namespace, acl.NamespaceCapabilityReadFS)
		logs := aclObj.AllowNsOp(args.QueryOptions.Namespace, acl.NamespaceCapabilityReadLogs)
		if !readfs && !logs {
			f.handleStreamResultError(structs.ErrPermissionDenied, nil, encoder)
			return
		}
	}

	// Verify the arguments.
	if args.AllocID == "" {
		f.handleStreamResultError(errors.New("missing AllocID"), helper.Int64ToPtr(400), encoder)
		return
	}

	// Retrieve the allocation
	snap, err := f.srv.State().Snapshot()
	if err != nil {
		f.handleStreamResultError(err, nil, encoder)
		return
	}

	alloc, err := snap.AllocByID(nil, args.AllocID)
	if err != nil {
		f.handleStreamResultError(err, nil, encoder)
		return
	}
	if alloc == nil {
		f.handleStreamResultError(fmt.Errorf("unknown alloc ID %q", args.AllocID), helper.Int64ToPtr(404), encoder)
		return
	}
	nodeID := alloc.NodeID

	// Get the connection to the client either by forwarding to another server
	// or creating a direct stream
	var clientConn net.Conn
	state, ok := f.srv.getNodeConn(nodeID)
	if !ok {
		// Determine the Server that has a connection to the node.
		srv, err := f.srv.serverWithNodeConn(nodeID, f.srv.Region())
		if err != nil {
			f.handleStreamResultError(err, nil, encoder)
			return
		}

		// Get a connection to the server
		conn, err := f.srv.streamingRpc(srv, "FileSystem.Logs")
		if err != nil {
			f.handleStreamResultError(err, nil, encoder)
			return
		}

		clientConn = conn
	} else {
		stream, err := NodeStreamingRpc(state.Session, "FileSystem.Logs")
		if err != nil {
			f.handleStreamResultError(err, nil, encoder)
			return
		}
		clientConn = stream
	}
	defer clientConn.Close()

	// Send the request.
	outEncoder := codec.NewEncoder(clientConn, structs.MsgpackHandle)
	if err := outEncoder.Encode(args); err != nil {
		f.handleStreamResultError(err, nil, encoder)
		return
	}

	structs.Bridge(conn, clientConn)
	return
}
<|endoftext|>


<commit_before>package demo_server

import (
	"html/template"
	"log"
	"net/http"
)

func aboutPage(ctx server, w http.ResponseWriter, req *http.Request) {
	t, err := template.ParseFiles(
		"../assets/base.templ.html",
		"../assets/about.templ.html",
	)
	if err != nil {
		log.Fatal(err)
	}

	if err := t.Execute(w, nil); err != nil {
		log.Fatal(err)
	}
}

func aboutPartial(ctx server, w http.ResponseWriter, req *http.Request) {
	t, err := template.ParseFiles(
		"../assets/about.templ.html",
	)
	if err != nil {
		log.Fatal(err)
	}

	if err := t.Execute(w, nil); err != nil {
		log.Fatal(err)
	}
}

func aboutUserPage(ctx server, w http.ResponseWriter, req *http.Request) {
	t, err := template.ParseFiles(
		"../assets/base.templ.html",
		"../assets/about.templ.html",
		"../assets/about-user.templ.html",
	)
	if err != nil {
		log.Fatal(err)
	}

	if err := t.Execute(w, nil); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>change template asset url for the moment<commit_after>package demo_server

import (
	"html/template"
	"log"
	"net/http"
)

func aboutPage(ctx server, w http.ResponseWriter, req *http.Request) {
	t, err := template.ParseFiles(
		"./assets/base.templ.html",
		"./assets/about.templ.html",
	)
	if err != nil {
		log.Fatal(err)
	}

	if err := t.Execute(w, nil); err != nil {
		log.Fatal(err)
	}
}

func aboutPartial(ctx server, w http.ResponseWriter, req *http.Request) {
	t, err := template.ParseFiles(
		"./assets/about.templ.html",
	)
	if err != nil {
		log.Fatal(err)
	}

	if err := t.Execute(w, nil); err != nil {
		log.Fatal(err)
	}
}

func aboutUserPage(ctx server, w http.ResponseWriter, req *http.Request) {
	t, err := template.ParseFiles(
		"./assets/base.templ.html",
		"./assets/about.templ.html",
		"./assets/about-user.templ.html",
	)
	if err != nil {
		log.Fatal(err)
	}

	if err := t.Execute(w, nil); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package stone

func lastWeightStones(stones []int) int {
	return 0
}
<commit_msg>solve 1046 by sorting<commit_after>package stone

import "sort"

func lastWeightStones(stones []int) int {
	return useSort(stones)
}

// useSort time complexity O(N^2lgN), space complexity O(1)
func useSort(stones []int) int {
	n := len(stones)
	sort.Ints(stones)
	i := n - 2
	for i >= 0 && i+1 < len(stones) {
		if stones[i] == stones[i+1] {
			stones = stones[:i]
			i -= 2
		} else {
			stones[i] = stones[i+1] - stones[i]
			stones = stones[:i+1]
			sort.Ints(stones)
			i--
		}
	}
	if len(stones) > 0 {
		return stones[0]
	}
	return 0
}
<|endoftext|>


<commit_before>// +build windows

package files

import (
	"testing"
)

var isSymlinkDataWin = []maybeln{
	{"", false},
	{"   ", false},
	{"?|!", false},
	{".notfound", false},
	{".", false},
	{"testdata", false},
	{"testdata/", false},
	{"testdata/files", false},
	{"testdata/files/", false},
	{"testdata/files/01.txt", false},
	{"testdata/files/01.txt/foo", false},
	{"testdata/files/linkto01", false},
	{"testdata/files/sub", false},
	{"testdata/files/sub/", false},
}

func TestIsSymlinkWin(t *testing.T) {
	for _, data := range isSymlinkDataWin {
		is := IsSymlink(data.path)
		if is != data.isln {
			t.Errorf(`Expected IsSymlink=%t for path "%s"`, data.isln, data.path)
		}
	}
}
<commit_msg>fix win test<commit_after>// +build windows

package files

import (
	"testing"
)

var isSymlinkDataWin = []maybeln{
	{"", false},
	{"   ", false},
	{"?|!", false},
	{".notfound", false},
	{".", false},
	{"testdata", false},
	{"testdata/", false},
	{"testdata/files", false},
	{"testdata/files/", false},
	{"testdata/files/01.txt", false},
	{"testdata/files/01.txt/foo", false},
	{"testdata/files/linkto01", true},
	{"testdata/files/sub", false},
	{"testdata/files/sub/", false},
}

func TestIsSymlinkWin(t *testing.T) {
	for _, data := range isSymlinkDataWin {
		is := IsSymlink(data.path)
		if is != data.isln {
			t.Errorf(`Expected IsSymlink=%t for path "%s"`, data.isln, data.path)
		}
	}
}
<|endoftext|>


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

package bgzf

import (
	"bufio"
	"bytes"
	"compress/flate"
	"compress/gzip"
	"io"
	"sync"
)

// countReader wraps flate.Reader, adding support for querying current offset.
type countReader struct {
	// Underlying Reader.
	fr flate.Reader

	// Offset within the underlying reader.
	off int64
}

// newCountReader returns a new countReader.
func newCountReader(r io.Reader) *countReader {
	switch r := r.(type) {
	case *countReader:
		panic("bgzf: illegal use of internal type")
	case flate.Reader:
		return &countReader{fr: r}
	default:
		return &countReader{fr: bufio.NewReader(r)}
	}
}

// Read is required to satisfy flate.Reader.
func (r *countReader) Read(p []byte) (int, error) {
	n, err := r.fr.Read(p)
	r.off += int64(n)
	return n, err
}

// ReadByte is required to satisfy flate.Reader.
func (r *countReader) ReadByte() (byte, error) {
	b, err := r.fr.ReadByte()
	if err == nil {
		r.off++
	}
	return b, err
}

// offset returns the current offset in the underlying reader.
func (r *countReader) offset() int64 { return r.off }

// seek moves the countReader to the specified offset using rs as the
// underlying reader.
func (r *countReader) seek(rs io.ReadSeeker, off int64) error {
	_, err := rs.Seek(off, 0)
	if err != nil {
		return err
	}

	type reseter interface {
		Reset(io.Reader)
	}
	switch cr := r.fr.(type) {
	case reseter:
		cr.Reset(rs)
	default:
		r.fr = newCountReader(rs)
	}
	r.off = off

	return nil
}

// buffer is a flate.Reader used by a decompressor to store read-ahead data.
type buffer struct {
	// Buffered compressed data from read ahead.
	off  int // Current position in buffered data.
	size int // Total size of buffered data.
	data [MaxBlockSize]byte
}

// Read provides the flate.Decompressor Read method.
func (r *buffer) Read(b []byte) (int, error) {
	if r.off >= r.size {
		return 0, io.EOF
	}
	if n := r.size - r.off; len(b) > n {
		b = b[:n]
	}
	n := copy(b, r.data[r.off:])
	r.off += n
	return n, nil
}

// ReadByte provides the flate.Decompressor ReadByte method.
func (r *buffer) ReadByte() (byte, error) {
	if r.off == r.size {
		return 0, io.EOF
	}
	b := r.data[r.off]
	r.off++
	return b, nil
}

// reset makes the buffer available to store data.
func (r *buffer) reset() { r.size = 0 }

// hasData returns whether the buffer has any data buffered.
func (r *buffer) hasData() bool { return r.size != 0 }

// readLimited reads n bytes into the buffer from the given source.
func (r *buffer) readLimited(n int, src *countReader) error {
	if r.hasData() {
		panic("bgzf: read into non-empty buffer")
	}
	r.off = 0
	var err error
	r.size, err = io.ReadFull(src, r.data[:n])
	return err
}

// decompressor is a gzip member decompressor worker.
type decompressor struct {
	owner *Reader

	gz gzip.Reader

	cr *countReader

	// Current block size.
	blockSize int

	// Buffered compressed data from read ahead.
	buf buffer

	// Decompressed data.
	wg  sync.WaitGroup
	blk Block

	err error
}

// Read provides the Read method for the decompressor's gzip.Reader.
func (d *decompressor) Read(b []byte) (int, error) {
	if d.buf.hasData() {
		return d.buf.Read(b)
	}
	return d.cr.Read(b)
}

// ReadByte provides the ReadByte method for the decompressor's gzip.Reader.
func (d *decompressor) ReadByte() (byte, error) {
	if d.buf.hasData() {
		return d.buf.ReadByte()
	}
	return d.cr.ReadByte()
}

// lazyBlock conditionally creates a ready to use Block.
func (d *decompressor) lazyBlock() {
	if d.blk == nil {
		if w, ok := d.owner.Cache.(Wrapper); ok {
			d.blk = w.Wrap(&block{owner: d.owner})
		} else {
			d.blk = &block{owner: d.owner}
		}
		return
	}
	if !d.blk.ownedBy(d.owner) {
		d.blk.setOwner(d.owner)
	}
}

// acquireHead gains the read head from the decompressor's owner.
func (d *decompressor) acquireHead() {
	d.wg.Add(1)
	d.cr = <-d.owner.head
}

// releaseHead releases the read head back to the decompressor's owner.
func (d *decompressor) releaseHead() {
	d.owner.head <- d.cr
	d.cr = nil // Defensively zero the reader.
}

// wait waits for the current member to be decompressed or fail, and returns
// the resulting error state.
func (d *decompressor) wait() (Block, error) {
	d.wg.Wait()
	blk := d.blk
	d.blk = nil
	return blk, d.err
}

func (d *decompressor) using(b Block) *decompressor { d.blk = b; return d }

// nextBlockAt makes the decompressor ready for reading decompressed data
// from its Block. It checks if there is a cached Block for the nextBase,
// otherwise it seeks to the correct location if decompressor is not
// correctly positioned, and then reads the compressed data and fills
// the decompressed Block.
// After nextBlockAt returns without error, the decompressor's Block
// holds a valid gzip.Header and base offset.
func (d *decompressor) nextBlockAt(off int64, rs io.ReadSeeker) *decompressor {
	blk, err := d.owner.cachedBlockFor(off)
	if err != nil {
		d.err = err
		return d
	}
	if blk != nil {
		// TODO(kortschak): Under some conditions, e.g. FIFO
		// cache we will be discarding a non-nil evicted Block.
		// Consider retaining these in a sync.Pool.
		d.owner.cachePut(d.blk)
		d.blk = blk
		return d
	}
	var retained bool
	d.blk, retained = d.owner.cachePut(d.blk)
	if retained {
		d.blk = nil
	}

	d.lazyBlock()

	d.acquireHead()
	defer d.releaseHead()

	if d.cr.offset() != off {
		// It should not be possible for the expected next block base
		// to be out of register with the count reader unless Seek
		// has been called, so we know the base reader must be an
		// io.ReadSeeker.
		if rs == nil {
			rs = d.owner.r.(io.ReadSeeker)
		}
		d.err = d.cr.seek(rs, off)
		if d.err != nil {
			d.wg.Done()
			return d
		}
	}

	d.blk.setBase(d.cr.offset())
	d.err = d.readMember()
	if d.err != nil {
		d.wg.Done()
		return d
	}
	d.blk.setHeader(d.gz.Header)

	// Decompress data into the decompressor's Block.
	go func() {
		d.err = d.blk.readFrom(&d.gz)
		d.wg.Done()
	}()

	return d
}

// expectedMemberSize returns the size of the BGZF conformant gzip member.
// It returns -1 if no BGZF block size field is found.
func expectedMemberSize(h gzip.Header) int {
	i := bytes.Index(h.Extra, bgzfExtraPrefix)
	if i < 0 || i+5 >= len(h.Extra) {
		return -1
	}
	return (int(h.Extra[i+4]) | int(h.Extra[i+5])<<8) + 1
}

// readMember buffers the gzip member starting the current decompressor offset.
func (d *decompressor) readMember() error {
	// Set the decompressor to Read from the underlying flate.Reader
	// and mark the starting offset from which the underlying reader
	// was used.
	d.buf.reset()
	mark := d.cr.offset()

	err := d.gz.Reset(d)
	if err != nil {
		d.blockSize = -1
		return err
	}

	d.blockSize = expectedMemberSize(d.gz.Header)
	if d.blockSize < 0 {
		return ErrNoBlockSize
	}

	// Read compressed data into the decompressor buffer until the
	// underlying flate.Reader is positioned at the end of the gzip
	// member in which the readMember call was made.
	return d.buf.readLimited(d.blockSize-int(d.cr.offset()-mark), d.cr)
}

// Offset is a BGZF virtual offset.
type Offset struct {
	File  int64
	Block uint16
}

// Chunk is a region of a BGZF file.
type Chunk struct {
	Begin Offset
	End   Offset
}

// Reader implements BGZF blocked gzip decompression.
type Reader struct {
	gzip.Header
	r io.Reader

	// head serialises access to the underlying
	// io.Reader.
	head chan *countReader

	// lastChunk is the virtual file offset
	// interval of the last successful read
	// or seek operation.
	lastChunk Chunk

	dec *decompressor

	current Block

	cacheLock sync.Mutex
	// Cache is the Reader block cache. If Cache is not nil,
	// the cache is queried for blocks before an attempt to
	// read from the underlying io.Reader.
	Cache Cache

	err error
}

// NewReader returns a new BGZF reader.
//
// The number of concurrent read decompressors is specified by
// rd (currently ignored).
func NewReader(r io.Reader, rd int) (*Reader, error) {
	bg := &Reader{
		r: r,

		head: make(chan *countReader, 1),
	}
	bg.head <- newCountReader(r)

	// Read the first block now so we can fail before
	// the first Read call if there is a problem.
	bg.dec = &decompressor{owner: bg}
	blk, err := bg.dec.nextBlockAt(0, nil).wait()
	if err != nil {
		return nil, err
	}
	bg.current = blk
	bg.Header = bg.current.header()

	return bg, nil
}

// Seek performs a seek operation to the given virtual offset.
func (bg *Reader) Seek(off Offset) error {
	rs, ok := bg.r.(io.ReadSeeker)
	if !ok {
		return ErrNotASeeker
	}

	if off.File != bg.current.Base() || !bg.current.hasData() {
		bg.current, bg.err = bg.dec.
			using(bg.current).
			nextBlockAt(off.File, rs).
			wait()
		bg.Header = bg.current.header()
		if bg.err != nil {
			return bg.err
		}
	}

	bg.err = bg.current.seek(int64(off.Block))
	if bg.err == nil {
		bg.lastChunk = Chunk{Begin: off, End: off}
	}

	return bg.err
}

// LastChunk returns the region of the BGZF file read by the last read
// operation or the resulting virtual offset of the last successful
// seek operation.
func (bg *Reader) LastChunk() Chunk { return bg.lastChunk }

// Close closes the reader and releases resources.
func (bg *Reader) Close() error {
	bg.Cache = nil
	if bg.err == io.EOF {
		return nil
	}
	return bg.err
}

// Read implements the io.Reader interface.
func (bg *Reader) Read(p []byte) (int, error) {
	if bg.err != nil {
		return 0, bg.err
	}

	// Discard leading empty blocks. This is an indexing
	// optimisation to avoid retaining useless members
	// in a BAI/CSI.
	for bg.current.len() == 0 {
		bg.current, bg.err = bg.nextBlock()
		bg.Header = bg.current.header()
		if bg.err != nil {
			return 0, bg.err
		}
	}

	bg.lastChunk.Begin = bg.current.txOffset()

	var n int
	for n < len(p) && bg.err == nil {
		var _n int
		_n, bg.err = bg.current.Read(p[n:])
		if _n > 0 {
			bg.lastChunk.End = bg.current.txOffset()
		}
		n += _n
		if bg.err == io.EOF {
			if n == len(p) {
				bg.err = nil
				break
			}

			bg.current, bg.err = bg.nextBlock()
			bg.Header = bg.current.header()
			if bg.err != nil {
				break
			}
		}
	}

	return n, bg.err
}

func (bg *Reader) nextBlock() (Block, error) {
	return bg.dec.
		using(bg.current).
		nextBlockAt(bg.current.nextBase(), nil).
		wait()
}

// cachedBlockFor returns true if the Reader has access to a cache
// and that cache holds the block with the given base and the correct
// owner, otherwise it returns nil. If the Block's owner is not correct,
// or the Block cannot seek to the start of its data, a non-nil error
// is returned.
func (bg *Reader) cachedBlockFor(base int64) (Block, error) {
	if bg.Cache == nil {
		return nil, nil
	}
	bg.cacheLock.Lock()
	defer bg.cacheLock.Unlock()
	blk := bg.Cache.Get(base)
	if blk != nil {
		if !blk.ownedBy(bg) {
			return nil, ErrContaminatedCache
		}
		err := blk.seek(0)
		if err != nil {
			return nil, err
		}
	}
	return blk, nil
}

// cachePut puts the given Block into the cache if it exists, it returns
// the Block that was evicted or b if it was not retained, and whether
// the Block was retained by the cache.
func (bg *Reader) cachePut(b Block) (evicted Block, retained bool) {
	if bg.Cache == nil || b == nil || !b.hasData() {
		return b, false
	}
	bg.cacheLock.Lock()
	defer bg.cacheLock.Unlock()
	return bg.Cache.Put(b)
}
<commit_msg>Move comment to relevant line<commit_after>// Copyright ©2012 The bíogo.bam Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package bgzf

import (
	"bufio"
	"bytes"
	"compress/flate"
	"compress/gzip"
	"io"
	"sync"
)

// countReader wraps flate.Reader, adding support for querying current offset.
type countReader struct {
	// Underlying Reader.
	fr flate.Reader

	// Offset within the underlying reader.
	off int64
}

// newCountReader returns a new countReader.
func newCountReader(r io.Reader) *countReader {
	switch r := r.(type) {
	case *countReader:
		panic("bgzf: illegal use of internal type")
	case flate.Reader:
		return &countReader{fr: r}
	default:
		return &countReader{fr: bufio.NewReader(r)}
	}
}

// Read is required to satisfy flate.Reader.
func (r *countReader) Read(p []byte) (int, error) {
	n, err := r.fr.Read(p)
	r.off += int64(n)
	return n, err
}

// ReadByte is required to satisfy flate.Reader.
func (r *countReader) ReadByte() (byte, error) {
	b, err := r.fr.ReadByte()
	if err == nil {
		r.off++
	}
	return b, err
}

// offset returns the current offset in the underlying reader.
func (r *countReader) offset() int64 { return r.off }

// seek moves the countReader to the specified offset using rs as the
// underlying reader.
func (r *countReader) seek(rs io.ReadSeeker, off int64) error {
	_, err := rs.Seek(off, 0)
	if err != nil {
		return err
	}

	type reseter interface {
		Reset(io.Reader)
	}
	switch cr := r.fr.(type) {
	case reseter:
		cr.Reset(rs)
	default:
		r.fr = newCountReader(rs)
	}
	r.off = off

	return nil
}

// buffer is a flate.Reader used by a decompressor to store read-ahead data.
type buffer struct {
	// Buffered compressed data from read ahead.
	off  int // Current position in buffered data.
	size int // Total size of buffered data.
	data [MaxBlockSize]byte
}

// Read provides the flate.Decompressor Read method.
func (r *buffer) Read(b []byte) (int, error) {
	if r.off >= r.size {
		return 0, io.EOF
	}
	if n := r.size - r.off; len(b) > n {
		b = b[:n]
	}
	n := copy(b, r.data[r.off:])
	r.off += n
	return n, nil
}

// ReadByte provides the flate.Decompressor ReadByte method.
func (r *buffer) ReadByte() (byte, error) {
	if r.off == r.size {
		return 0, io.EOF
	}
	b := r.data[r.off]
	r.off++
	return b, nil
}

// reset makes the buffer available to store data.
func (r *buffer) reset() { r.size = 0 }

// hasData returns whether the buffer has any data buffered.
func (r *buffer) hasData() bool { return r.size != 0 }

// readLimited reads n bytes into the buffer from the given source.
func (r *buffer) readLimited(n int, src *countReader) error {
	if r.hasData() {
		panic("bgzf: read into non-empty buffer")
	}
	r.off = 0
	var err error
	r.size, err = io.ReadFull(src, r.data[:n])
	return err
}

// decompressor is a gzip member decompressor worker.
type decompressor struct {
	owner *Reader

	gz gzip.Reader

	cr *countReader

	// Current block size.
	blockSize int

	// Buffered compressed data from read ahead.
	buf buffer

	// Decompressed data.
	wg  sync.WaitGroup
	blk Block

	err error
}

// Read provides the Read method for the decompressor's gzip.Reader.
func (d *decompressor) Read(b []byte) (int, error) {
	if d.buf.hasData() {
		return d.buf.Read(b)
	}
	return d.cr.Read(b)
}

// ReadByte provides the ReadByte method for the decompressor's gzip.Reader.
func (d *decompressor) ReadByte() (byte, error) {
	if d.buf.hasData() {
		return d.buf.ReadByte()
	}
	return d.cr.ReadByte()
}

// lazyBlock conditionally creates a ready to use Block.
func (d *decompressor) lazyBlock() {
	if d.blk == nil {
		if w, ok := d.owner.Cache.(Wrapper); ok {
			d.blk = w.Wrap(&block{owner: d.owner})
		} else {
			d.blk = &block{owner: d.owner}
		}
		return
	}
	if !d.blk.ownedBy(d.owner) {
		d.blk.setOwner(d.owner)
	}
}

// acquireHead gains the read head from the decompressor's owner.
func (d *decompressor) acquireHead() {
	d.wg.Add(1)
	d.cr = <-d.owner.head
}

// releaseHead releases the read head back to the decompressor's owner.
func (d *decompressor) releaseHead() {
	d.owner.head <- d.cr
	d.cr = nil // Defensively zero the reader.
}

// wait waits for the current member to be decompressed or fail, and returns
// the resulting error state.
func (d *decompressor) wait() (Block, error) {
	d.wg.Wait()
	blk := d.blk
	d.blk = nil
	return blk, d.err
}

func (d *decompressor) using(b Block) *decompressor { d.blk = b; return d }

// nextBlockAt makes the decompressor ready for reading decompressed data
// from its Block. It checks if there is a cached Block for the nextBase,
// otherwise it seeks to the correct location if decompressor is not
// correctly positioned, and then reads the compressed data and fills
// the decompressed Block.
// After nextBlockAt returns without error, the decompressor's Block
// holds a valid gzip.Header and base offset.
func (d *decompressor) nextBlockAt(off int64, rs io.ReadSeeker) *decompressor {
	blk, err := d.owner.cachedBlockFor(off)
	if err != nil {
		d.err = err
		return d
	}
	if blk != nil {
		// TODO(kortschak): Under some conditions, e.g. FIFO
		// cache we will be discarding a non-nil evicted Block.
		// Consider retaining these in a sync.Pool.
		d.owner.cachePut(d.blk)
		d.blk = blk
		return d
	}
	var retained bool
	d.blk, retained = d.owner.cachePut(d.blk)
	if retained {
		d.blk = nil
	}

	d.lazyBlock()

	d.acquireHead()
	defer d.releaseHead()

	if d.cr.offset() != off {
		if rs == nil {
			// It should not be possible for the expected next block base
			// to be out of register with the count reader unless Seek
			// has been called, so we know the base reader must be an
			// io.ReadSeeker.
			rs = d.owner.r.(io.ReadSeeker)
		}
		d.err = d.cr.seek(rs, off)
		if d.err != nil {
			d.wg.Done()
			return d
		}
	}

	d.blk.setBase(d.cr.offset())
	d.err = d.readMember()
	if d.err != nil {
		d.wg.Done()
		return d
	}
	d.blk.setHeader(d.gz.Header)

	// Decompress data into the decompressor's Block.
	go func() {
		d.err = d.blk.readFrom(&d.gz)
		d.wg.Done()
	}()

	return d
}

// expectedMemberSize returns the size of the BGZF conformant gzip member.
// It returns -1 if no BGZF block size field is found.
func expectedMemberSize(h gzip.Header) int {
	i := bytes.Index(h.Extra, bgzfExtraPrefix)
	if i < 0 || i+5 >= len(h.Extra) {
		return -1
	}
	return (int(h.Extra[i+4]) | int(h.Extra[i+5])<<8) + 1
}

// readMember buffers the gzip member starting the current decompressor offset.
func (d *decompressor) readMember() error {
	// Set the decompressor to Read from the underlying flate.Reader
	// and mark the starting offset from which the underlying reader
	// was used.
	d.buf.reset()
	mark := d.cr.offset()

	err := d.gz.Reset(d)
	if err != nil {
		d.blockSize = -1
		return err
	}

	d.blockSize = expectedMemberSize(d.gz.Header)
	if d.blockSize < 0 {
		return ErrNoBlockSize
	}

	// Read compressed data into the decompressor buffer until the
	// underlying flate.Reader is positioned at the end of the gzip
	// member in which the readMember call was made.
	return d.buf.readLimited(d.blockSize-int(d.cr.offset()-mark), d.cr)
}

// Offset is a BGZF virtual offset.
type Offset struct {
	File  int64
	Block uint16
}

// Chunk is a region of a BGZF file.
type Chunk struct {
	Begin Offset
	End   Offset
}

// Reader implements BGZF blocked gzip decompression.
type Reader struct {
	gzip.Header
	r io.Reader

	// head serialises access to the underlying
	// io.Reader.
	head chan *countReader

	// lastChunk is the virtual file offset
	// interval of the last successful read
	// or seek operation.
	lastChunk Chunk

	dec *decompressor

	current Block

	cacheLock sync.Mutex
	// Cache is the Reader block cache. If Cache is not nil,
	// the cache is queried for blocks before an attempt to
	// read from the underlying io.Reader.
	Cache Cache

	err error
}

// NewReader returns a new BGZF reader.
//
// The number of concurrent read decompressors is specified by
// rd (currently ignored).
func NewReader(r io.Reader, rd int) (*Reader, error) {
	bg := &Reader{
		r: r,

		head: make(chan *countReader, 1),
	}
	bg.head <- newCountReader(r)

	// Read the first block now so we can fail before
	// the first Read call if there is a problem.
	bg.dec = &decompressor{owner: bg}
	blk, err := bg.dec.nextBlockAt(0, nil).wait()
	if err != nil {
		return nil, err
	}
	bg.current = blk
	bg.Header = bg.current.header()

	return bg, nil
}

// Seek performs a seek operation to the given virtual offset.
func (bg *Reader) Seek(off Offset) error {
	rs, ok := bg.r.(io.ReadSeeker)
	if !ok {
		return ErrNotASeeker
	}

	if off.File != bg.current.Base() || !bg.current.hasData() {
		bg.current, bg.err = bg.dec.
			using(bg.current).
			nextBlockAt(off.File, rs).
			wait()
		bg.Header = bg.current.header()
		if bg.err != nil {
			return bg.err
		}
	}

	bg.err = bg.current.seek(int64(off.Block))
	if bg.err == nil {
		bg.lastChunk = Chunk{Begin: off, End: off}
	}

	return bg.err
}

// LastChunk returns the region of the BGZF file read by the last read
// operation or the resulting virtual offset of the last successful
// seek operation.
func (bg *Reader) LastChunk() Chunk { return bg.lastChunk }

// Close closes the reader and releases resources.
func (bg *Reader) Close() error {
	bg.Cache = nil
	if bg.err == io.EOF {
		return nil
	}
	return bg.err
}

// Read implements the io.Reader interface.
func (bg *Reader) Read(p []byte) (int, error) {
	if bg.err != nil {
		return 0, bg.err
	}

	// Discard leading empty blocks. This is an indexing
	// optimisation to avoid retaining useless members
	// in a BAI/CSI.
	for bg.current.len() == 0 {
		bg.current, bg.err = bg.nextBlock()
		bg.Header = bg.current.header()
		if bg.err != nil {
			return 0, bg.err
		}
	}

	bg.lastChunk.Begin = bg.current.txOffset()

	var n int
	for n < len(p) && bg.err == nil {
		var _n int
		_n, bg.err = bg.current.Read(p[n:])
		if _n > 0 {
			bg.lastChunk.End = bg.current.txOffset()
		}
		n += _n
		if bg.err == io.EOF {
			if n == len(p) {
				bg.err = nil
				break
			}

			bg.current, bg.err = bg.nextBlock()
			bg.Header = bg.current.header()
			if bg.err != nil {
				break
			}
		}
	}

	return n, bg.err
}

func (bg *Reader) nextBlock() (Block, error) {
	return bg.dec.
		using(bg.current).
		nextBlockAt(bg.current.nextBase(), nil).
		wait()
}

// cachedBlockFor returns true if the Reader has access to a cache
// and that cache holds the block with the given base and the correct
// owner, otherwise it returns nil. If the Block's owner is not correct,
// or the Block cannot seek to the start of its data, a non-nil error
// is returned.
func (bg *Reader) cachedBlockFor(base int64) (Block, error) {
	if bg.Cache == nil {
		return nil, nil
	}
	bg.cacheLock.Lock()
	defer bg.cacheLock.Unlock()
	blk := bg.Cache.Get(base)
	if blk != nil {
		if !blk.ownedBy(bg) {
			return nil, ErrContaminatedCache
		}
		err := blk.seek(0)
		if err != nil {
			return nil, err
		}
	}
	return blk, nil
}

// cachePut puts the given Block into the cache if it exists, it returns
// the Block that was evicted or b if it was not retained, and whether
// the Block was retained by the cache.
func (bg *Reader) cachePut(b Block) (evicted Block, retained bool) {
	if bg.Cache == nil || b == nil || !b.hasData() {
		return b, false
	}
	bg.cacheLock.Lock()
	defer bg.cacheLock.Unlock()
	return bg.Cache.Put(b)
}
<|endoftext|>


<commit_before>// Basic package for faking Stdio.
//
// Eli Bendersky [https://eli.thegreenplace.net]
// This code is in the public domain.
package fakestdio

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
)

// FakeStdio can be used to fake stdin and capture stdout.
// Between creating a new FakeStdio and calling ReadAndRestore on it,
// code reading os.Stdin will get the contents of stdinText passed to New.
// Output to os.Stdout will be captured and returned from ReadAndRestore.
// FakeStdio is not reusable; don't attempt to use it after calling
// ReadAndRestore, but it should be safe to create a new FakeStdio.
type FakeStdio struct {
	origStdout   *os.File
	stdoutReader *os.File

	outCh chan []byte

	origStdin   *os.File
	stdinWriter *os.File
}

func New(stdinText string) (*FakeStdio, error) {
	// Pipe for stdin.
	//
	//                 ======
	//  w ------------->||||------> r
	// (stdinWriter)   ======      (os.Stdin)
	stdinReader, stdinWriter, err := os.Pipe()
	if err != nil {
		return nil, err
	}

	// Pipe for stdout.
	//
	//               ======
	//  w ----------->||||------> r
	// (os.Stdout)   ======      (stdoutReader)
	stdoutReader, stdoutWriter, err := os.Pipe()
	if err != nil {
		return nil, err
	}

	origStdin := os.Stdin
	os.Stdin = stdinReader

	_, err = stdinWriter.Write([]byte(stdinText))
	if err != nil {
		stdinWriter.Close()
		os.Stdin = origStdin
		return nil, err
	}

	origStdout := os.Stdout
	os.Stdout = stdoutWriter

	outCh := make(chan []byte)

	// This goroutine reads stdout into a buffer in the background.
	go func() {
		var b bytes.Buffer
		if _, err := io.Copy(&b, stdoutReader); err != nil {
			log.Println(err)
		}
		outCh <- b.Bytes()
	}()

	return &FakeStdio{
		origStdout:   origStdout,
		stdoutReader: stdoutReader,
		outCh:        outCh,
		origStdin:    origStdin,
		stdinWriter:  stdinWriter,
	}, nil
}

// CloseStdin closes the fake stdin. This may be necessary if the process has
// logic for reading stdin until EOF; otherwise such code would block forever.
func (sf *FakeStdio) CloseStdin() {
	if sf.stdinWriter != nil {
		sf.stdinWriter.Close()
		sf.stdinWriter = nil
	}
}

// ReadAndRestore collects all captured stdout and returns it; it also restores
// os.Stdin and os.Stdout to their original values.
func (sf *FakeStdio) ReadAndRestore() ([]byte, error) {
	if sf.stdoutReader == nil {
		return nil, fmt.Errorf("ReadAndRestore from closed FakeStdio")
	}

	// Close the writer side of the faked stdout pipe. This signals to the
	// background goroutine it that it should exit.
	os.Stdout.Close()
	out := <-sf.outCh

	os.Stdout = sf.origStdout
	os.Stdin = sf.origStdin

	if sf.stdoutReader != nil {
		sf.stdoutReader.Close()
		sf.stdoutReader = nil
	}

	if sf.stdinWriter != nil {
		sf.stdinWriter.Close()
		sf.stdinWriter = nil
	}

	return out, nil
}
<commit_msg>Fix typo in comment<commit_after>// Basic package for faking Stdio.
//
// Eli Bendersky [https://eli.thegreenplace.net]
// This code is in the public domain.
package fakestdio

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
)

// FakeStdio can be used to fake stdin and capture stdout.
// Between creating a new FakeStdio and calling ReadAndRestore on it,
// code reading os.Stdin will get the contents of stdinText passed to New.
// Output to os.Stdout will be captured and returned from ReadAndRestore.
// FakeStdio is not reusable; don't attempt to use it after calling
// ReadAndRestore, but it should be safe to create a new FakeStdio.
type FakeStdio struct {
	origStdout   *os.File
	stdoutReader *os.File

	outCh chan []byte

	origStdin   *os.File
	stdinWriter *os.File
}

func New(stdinText string) (*FakeStdio, error) {
	// Pipe for stdin.
	//
	//                 ======
	//  w ------------->||||------> r
	// (stdinWriter)   ======      (os.Stdin)
	stdinReader, stdinWriter, err := os.Pipe()
	if err != nil {
		return nil, err
	}

	// Pipe for stdout.
	//
	//               ======
	//  w ----------->||||------> r
	// (os.Stdout)   ======      (stdoutReader)
	stdoutReader, stdoutWriter, err := os.Pipe()
	if err != nil {
		return nil, err
	}

	origStdin := os.Stdin
	os.Stdin = stdinReader

	_, err = stdinWriter.Write([]byte(stdinText))
	if err != nil {
		stdinWriter.Close()
		os.Stdin = origStdin
		return nil, err
	}

	origStdout := os.Stdout
	os.Stdout = stdoutWriter

	outCh := make(chan []byte)

	// This goroutine reads stdout into a buffer in the background.
	go func() {
		var b bytes.Buffer
		if _, err := io.Copy(&b, stdoutReader); err != nil {
			log.Println(err)
		}
		outCh <- b.Bytes()
	}()

	return &FakeStdio{
		origStdout:   origStdout,
		stdoutReader: stdoutReader,
		outCh:        outCh,
		origStdin:    origStdin,
		stdinWriter:  stdinWriter,
	}, nil
}

// CloseStdin closes the fake stdin. This may be necessary if the process has
// logic for reading stdin until EOF; otherwise such code would block forever.
func (sf *FakeStdio) CloseStdin() {
	if sf.stdinWriter != nil {
		sf.stdinWriter.Close()
		sf.stdinWriter = nil
	}
}

// ReadAndRestore collects all captured stdout and returns it; it also restores
// os.Stdin and os.Stdout to their original values.
func (sf *FakeStdio) ReadAndRestore() ([]byte, error) {
	if sf.stdoutReader == nil {
		return nil, fmt.Errorf("ReadAndRestore from closed FakeStdio")
	}

	// Close the writer side of the faked stdout pipe. This signals to the
	// background goroutine that it should exit.
	os.Stdout.Close()
	out := <-sf.outCh

	os.Stdout = sf.origStdout
	os.Stdin = sf.origStdin

	if sf.stdoutReader != nil {
		sf.stdoutReader.Close()
		sf.stdoutReader = nil
	}

	if sf.stdinWriter != nil {
		sf.stdinWriter.Close()
		sf.stdinWriter = nil
	}

	return out, nil
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"koding/kontrol/kontrolhelper"
	"koding/tools/amqputil"
	"koding/tools/config"
	"koding/tools/lifecycle"
	"koding/tools/logger"
	"koding/tools/sockjs"
	"koding/tools/utils"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/streadway/amqp"
)

var (
	log                    = logger.New("broker")
	routeMap               = make(map[string]([]*sockjs.Session))
	socketSubscriptionsMap = make(map[string]*map[string]bool)
	globalMapMutex         sync.Mutex
)

type Broker struct {
	Hostname          string
	ServiceUniqueName string
	PublishConn       *amqp.Connection
}

func NewBroker() *Broker {
	// returns os.Hostname() if config.BrokerDomain is empty, otherwise it just
	// returns config.BrokerDomain back
	brokerHostname := kontrolhelper.CustomHostname(config.BrokerDomain)
	sanitizedHostname := strings.Replace(brokerHostname, ".", "_", -1)
	serviceUniqueName := "broker" + "|" + sanitizedHostname

	return &Broker{
		Hostname:          brokerHostname,
		ServiceUniqueName: serviceUniqueName,
	}
}

func main() {
	lifecycle.Startup("broker", false)
	logger.RunGaugesLoop(log)

	broker := NewBroker()
	broker.registerToKontrol()

	go broker.startSockJS()
	broker.startAMQP() // blocking

	time.Sleep(5 * time.Second) // give amqputil time to log connection error
}

func (b *Broker) registerToKontrol() {
	if err := kontrolhelper.RegisterToKontrol(
		"broker", // servicename
		"broker",
		b.ServiceUniqueName,
		config.Uuid,
		b.Hostname,
		config.Current.Broker.Port,
	); err != nil {
		panic(err)
	}
}

func (b *Broker) startAMQP() {
	b.PublishConn = amqputil.CreateConnection("broker")
	defer b.PublishConn.Close()

	consumeConn := amqputil.CreateConnection("broker")
	defer consumeConn.Close()

	consumeChannel := amqputil.CreateChannel(consumeConn)
	defer consumeChannel.Close()

	presenceQueue := amqputil.JoinPresenceExchange(
		consumeChannel,      // channel
		"services-presence", // exchange
		"broker",            // serviceType
		"broker",            // serviceGenericName
		b.ServiceUniqueName, // serviceUniqueName
		false,               // loadBalancing
	)

	go func() {
		sigusr1Channel := make(chan os.Signal)
		signal.Notify(sigusr1Channel, syscall.SIGUSR1)
		<-sigusr1Channel
		consumeChannel.QueueDelete(presenceQueue, false, false, false)
	}()

	stream := amqputil.DeclareBindConsumeQueue(consumeChannel, "topic", "broker", "#", false)

	if err := consumeChannel.ExchangeDeclare(
		"updateInstances", // name
		"fanout",          // kind
		false,             // durable
		false,             // autoDelete
		false,             // internal
		false,             // noWait
		nil,               // args
	); err != nil {
		panic(err)
	}

	if err := consumeChannel.ExchangeBind("broker", "", "updateInstances", false, nil); err != nil {
		panic(err)
	}

	// start to listen from "broker" topic exchange
	for amqpMessage := range stream {
		routingKey := amqpMessage.RoutingKey
		payload := json.RawMessage(utils.FilterInvalidUTF8(amqpMessage.Body))

		pos := strings.IndexRune(routingKey, '.') // skip first dot, since we want at least two components to always include the secret
		for pos != -1 && pos < len(routingKey) {
			index := strings.IndexRune(routingKey[pos+1:], '.')
			pos += index + 1
			if index == -1 {
				pos = len(routingKey)
			}
			prefix := routingKey[:pos]
			globalMapMutex.Lock()
			for _, routeSession := range routeMap[prefix] {
				sendToClient(routeSession, routingKey, &payload)
			}
			globalMapMutex.Unlock()
		}
	}
}

func (b *Broker) startSockJS() {
	service := sockjs.NewService(
		config.Current.Client.StaticFilesBaseUrl+"/js/sock.js",
		10*time.Minute,
		b.sockjsSession,
	)
	defer service.Close()

	service.MaxReceivedPerSecond = 50
	service.ErrorHandler = log.LogError

	server := &http.Server{
		Handler: &sockjs.Mux{
			Handlers: map[string]http.Handler{
				"/subscribe": service,
				"/buildnumber": http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					w.Header().Set("Content-Type", "text/plain")
					w.Write([]byte(strconv.Itoa(config.Current.BuildNumber)))
				}),
			},
		},
	}

	var listener net.Listener
	listener, err := net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP(config.Current.Broker.IP), Port: config.Current.Broker.Port})
	if err != nil {
		log.Fatal(err)
	}

	if config.Current.Broker.CertFile != "" {
		cert, err := tls.LoadX509KeyPair(config.Current.Broker.CertFile, config.Current.Broker.KeyFile)
		if err != nil {
			log.Fatal(err)
		}
		listener = tls.NewListener(listener, &tls.Config{
			NextProtos:   []string{"http/1.1"},
			Certificates: []tls.Certificate{cert},
		})
	}

	lastErrorTime := time.Now()
	for {
		err := server.Serve(listener)
		if err != nil {
			log.Warning("Server error: %v", err)
			if time.Now().Sub(lastErrorTime) < time.Second {
				log.Fatal(nil)
			}
			lastErrorTime = time.Now()
		}
	}

}

func sendToClient(session *sockjs.Session, routingKey string, payload interface{}) {
	var message struct {
		RoutingKey string      `json:"routingKey"`
		Payload    interface{} `json:"payload"`
	}
	message.RoutingKey = routingKey
	message.Payload = payload
	if !session.Send(message) {
		session.Close()
		log.Warning("Dropped session because of broker to client buffer overflow. %v", session.Tag)
	}
}

func (b *Broker) sockjsSession(session *sockjs.Session) {
	changeClientsGauge := lifecycle.CreateClientsGauge()
	changeNewClientsGauge := logger.CreateCounterGauge("newClients", logger.NoUnit, true)
	changeWebsocketClientsGauge := logger.CreateCounterGauge("websocketClients", logger.NoUnit, false)

	defer log.RecoverAndLog()

	r := make([]byte, 128/8)
	rand.Read(r)
	socketId := base64.StdEncoding.EncodeToString(r)
	session.Tag = socketId

	log.Debug("Client connected: %v", socketId)
	changeClientsGauge(1)
	changeNewClientsGauge(1)
	if session.IsWebsocket {
		changeWebsocketClientsGauge(1)
	}
	defer func() {
		log.Debug("Client disconnected: %v", socketId)
		changeClientsGauge(-1)
		if session.IsWebsocket {
			changeWebsocketClientsGauge(-1)
		}
	}()

	var controlChannel *amqp.Channel
	var lastPayload string
	resetControlChannel := func() {
		if controlChannel != nil {
			controlChannel.Close()
		}
		var err error
		controlChannel, err = b.PublishConn.Channel()
		if err != nil {
			panic(err)
		}
		go func() {
			defer log.RecoverAndLog()

			for amqpErr := range controlChannel.NotifyClose(make(chan *amqp.Error)) {
				if !(strings.Contains(amqpErr.Error(), "NOT_FOUND") && (strings.Contains(amqpErr.Error(), "koding-social-") || strings.Contains(amqpErr.Error(), "auth-"))) {
					log.Warning("AMQP channel: %v Last publish payload: %v", amqpErr.Error(), lastPayload)
				}

				sendToClient(session, "broker.error", map[string]interface{}{"code": amqpErr.Code, "reason": amqpErr.Reason, "server": amqpErr.Server, "recover": amqpErr.Recover})
			}
		}()
	}
	resetControlChannel()
	defer func() { controlChannel.Close() }()

	subscriptions := make(map[string]bool)
	globalMapMutex.Lock()
	socketSubscriptionsMap[socketId] = &subscriptions
	globalMapMutex.Unlock()

	removeFromRouteMap := func(routingKeyPrefix string) {
		routeSessions := routeMap[routingKeyPrefix]
		for i, routeSession := range routeSessions {
			if routeSession == session {
				routeSessions[i] = routeSessions[len(routeSessions)-1]
				routeSessions = routeSessions[:len(routeSessions)-1]
				break
			}
		}
		if len(routeSessions) == 0 {
			delete(routeMap, routingKeyPrefix)
			return
		}
		routeMap[routingKeyPrefix] = routeSessions
	}

	subscribe := func(routingKeyPrefix string) {
		if subscriptions[routingKeyPrefix] {
			log.Warning("Duplicate subscription to same routing key. %v %v", session.Tag, routingKeyPrefix)
			return
		}
		if len(subscriptions) > 0 && len(subscriptions)%2000 == 0 {
			log.Warning("Client with more than %v subscriptions %v", strconv.Itoa(len(subscriptions)), session.Tag)
		}
		routeMap[routingKeyPrefix] = append(routeMap[routingKeyPrefix], session)
		subscriptions[routingKeyPrefix] = true
	}

	unsubscribe := func(routingKeyPrefix string) {
		removeFromRouteMap(routingKeyPrefix)
		delete(subscriptions, routingKeyPrefix)
	}

	defer func() {
		globalMapMutex.Lock()
		for routingKeyPrefix := range subscriptions {
			removeFromRouteMap(routingKeyPrefix)
		}
		globalMapMutex.Unlock()

		time.AfterFunc(5*time.Minute, func() {
			globalMapMutex.Lock()
			delete(socketSubscriptionsMap, socketId)
			globalMapMutex.Unlock()
		})

		for {
			err := controlChannel.Publish(config.Current.Broker.AuthAllExchange, "broker.clientDisconnected", false, false, amqp.Publishing{Body: []byte(socketId)})
			if err == nil {
				break
			}
			if amqpError, isAmqpError := err.(*amqp.Error); !isAmqpError || amqpError.Code != 504 {
				panic(err)
			}
			resetControlChannel()
		}
	}()

	err := controlChannel.Publish(config.Current.Broker.AuthAllExchange, "broker.clientConnected", false, false, amqp.Publishing{Body: []byte(socketId)})
	if err != nil {
		panic(err)
	}

	sendToClient(session, "broker.connected", socketId)

	for data := range session.ReceiveChan {
		if data == nil || session.Closed {
			break
		}

		defer log.RecoverAndLog()

		message := data.(map[string]interface{})
		log.Debug("Received message: %v", message)

		action := message["action"]
		switch action {
		case "subscribe":
			globalMapMutex.Lock()
			defer globalMapMutex.Unlock()
			for _, routingKeyPrefix := range strings.Split(message["routingKeyPrefix"].(string), " ") {
				subscribe(routingKeyPrefix)
			}
			sendToClient(session, "broker.subscribed", message["routingKeyPrefix"])

		case "resubscribe":
			globalMapMutex.Lock()
			defer globalMapMutex.Unlock()
			oldSubscriptions, found := socketSubscriptionsMap[message["socketId"].(string)]
			if found {
				for routingKeyPrefix := range *oldSubscriptions {
					subscribe(routingKeyPrefix)
				}
			}
			sendToClient(session, "broker.resubscribed", found)

		case "unsubscribe":
			globalMapMutex.Lock()
			defer globalMapMutex.Unlock()
			for _, routingKeyPrefix := range strings.Split(message["routingKeyPrefix"].(string), " ") {
				unsubscribe(routingKeyPrefix)
			}

		case "publish":
			exchange := message["exchange"].(string)
			routingKey := message["routingKey"].(string)
			if !strings.HasPrefix(routingKey, "client.") {
				log.Warning("Invalid routing key: message: %v socketId: %v", message, socketId)
				return
			}
			for {
				lastPayload = ""
				err := controlChannel.Publish(exchange, routingKey, false, false, amqp.Publishing{CorrelationId: socketId, Body: []byte(message["payload"].(string))})
				if err == nil {
					lastPayload = message["payload"].(string)
					break
				}
				if amqpError, isAmqpError := err.(*amqp.Error); !isAmqpError || amqpError.Code != 504 {
					log.Warning("payload: %v routing key: %v exchange: %v err: %v", message["payload"], message["routingKey"], message["exchange"], err)
				}
				time.Sleep(time.Second / 4) // penalty for crashing the AMQP channel
				resetControlChannel()
			}

		case "ping":
			sendToClient(session, "broker.pong", nil)

		default:
			log.Warning("Invalid action. message: %v socketId: %v", message, socketId)

		}
	}
}
<commit_msg>broker: more refactorings, WIP<commit_after>package main

import (
	"crypto/rand"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"koding/kontrol/kontrolhelper"
	"koding/tools/amqputil"
	"koding/tools/config"
	"koding/tools/lifecycle"
	"koding/tools/logger"
	"koding/tools/sockjs"
	"koding/tools/utils"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/streadway/amqp"
)

var (
	log                    = logger.New("broker")
	routeMap               = make(map[string]([]*sockjs.Session))
	socketSubscriptionsMap = make(map[string]*map[string]bool)
	globalMapMutex         sync.Mutex
)

type Broker struct {
	Hostname          string
	ServiceUniqueName string
	PublishConn       *amqp.Connection
}

func NewBroker() *Broker {
	// returns os.Hostname() if config.BrokerDomain is empty, otherwise it just
	// returns config.BrokerDomain back
	brokerHostname := kontrolhelper.CustomHostname(config.BrokerDomain)
	sanitizedHostname := strings.Replace(brokerHostname, ".", "_", -1)
	serviceUniqueName := "broker" + "|" + sanitizedHostname

	return &Broker{
		Hostname:          brokerHostname,
		ServiceUniqueName: serviceUniqueName,
	}
}

func main() {
	lifecycle.Startup("broker", false)
	logger.RunGaugesLoop(log)

	broker := NewBroker()
	broker.registerToKontrol()

	go broker.startSockJS()
	broker.startAMQP() // blocking

	time.Sleep(5 * time.Second) // give amqputil time to log connection error
}

func (b *Broker) registerToKontrol() {
	if err := kontrolhelper.RegisterToKontrol(
		"broker", // servicename
		"broker",
		b.ServiceUniqueName,
		config.Uuid,
		b.Hostname,
		config.Current.Broker.Port,
	); err != nil {
		panic(err)
	}
}

func (b *Broker) startAMQP() {
	b.PublishConn = amqputil.CreateConnection("broker")
	defer b.PublishConn.Close()

	consumeConn := amqputil.CreateConnection("broker")
	defer consumeConn.Close()

	consumeChannel := amqputil.CreateChannel(consumeConn)
	defer consumeChannel.Close()

	presenceQueue := amqputil.JoinPresenceExchange(
		consumeChannel,      // channel
		"services-presence", // exchange
		"broker",            // serviceType
		"broker",            // serviceGenericName
		b.ServiceUniqueName, // serviceUniqueName
		false,               // loadBalancing
	)

	go func() {
		sigusr1Channel := make(chan os.Signal)
		signal.Notify(sigusr1Channel, syscall.SIGUSR1)
		<-sigusr1Channel
		consumeChannel.QueueDelete(presenceQueue, false, false, false)
	}()

	stream := amqputil.DeclareBindConsumeQueue(consumeChannel, "topic", "broker", "#", false)

	if err := consumeChannel.ExchangeDeclare(
		"updateInstances", // name
		"fanout",          // kind
		false,             // durable
		false,             // autoDelete
		false,             // internal
		false,             // noWait
		nil,               // args
	); err != nil {
		panic(err)
	}

	if err := consumeChannel.ExchangeBind("broker", "", "updateInstances", false, nil); err != nil {
		panic(err)
	}

	// start to listen from "broker" topic exchange
	for amqpMessage := range stream {
		routingKey := amqpMessage.RoutingKey
		payload := json.RawMessage(utils.FilterInvalidUTF8(amqpMessage.Body))

		pos := strings.IndexRune(routingKey, '.') // skip first dot, since we want at least two components to always include the secret
		for pos != -1 && pos < len(routingKey) {
			index := strings.IndexRune(routingKey[pos+1:], '.')
			pos += index + 1
			if index == -1 {
				pos = len(routingKey)
			}
			prefix := routingKey[:pos]
			globalMapMutex.Lock()
			for _, routeSession := range routeMap[prefix] {
				sendToClient(routeSession, routingKey, &payload)
			}
			globalMapMutex.Unlock()
		}
	}
}

func (b *Broker) startSockJS() {
	service := sockjs.NewService(
		config.Current.Client.StaticFilesBaseUrl+"/js/sock.js",
		10*time.Minute,
		b.sockjsSession,
	)
	defer service.Close()

	service.MaxReceivedPerSecond = 50
	service.ErrorHandler = log.LogError

	// TODO use http.Mux instead of sockjs.Mux.
	server := &http.Server{
		Handler: &sockjs.Mux{
			Handlers: map[string]http.Handler{
				"/subscribe": service,
				"/buildnumber": http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					w.Header().Set("Content-Type", "text/plain")
					w.Write([]byte(strconv.Itoa(config.Current.BuildNumber)))
				}),
			},
		},
	}

	var listener net.Listener
	listener, err := net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP(config.Current.Broker.IP), Port: config.Current.Broker.Port})
	if err != nil {
		log.Fatal(err)
	}

	if config.Current.Broker.CertFile != "" {
		cert, err := tls.LoadX509KeyPair(config.Current.Broker.CertFile, config.Current.Broker.KeyFile)
		if err != nil {
			log.Fatal(err)
		}
		listener = tls.NewListener(listener, &tls.Config{
			NextProtos:   []string{"http/1.1"},
			Certificates: []tls.Certificate{cert},
		})
	}

	lastErrorTime := time.Now()
	for {
		err := server.Serve(listener)
		if err != nil {
			log.Warning("Server error: %v", err)
			if time.Now().Sub(lastErrorTime) < time.Second {
				log.Fatal(nil)
			}
			lastErrorTime = time.Now()
		}
	}

}

func sendToClient(session *sockjs.Session, routingKey string, payload interface{}) {
	var message struct {
		RoutingKey string      `json:"routingKey"`
		Payload    interface{} `json:"payload"`
	}
	message.RoutingKey = routingKey
	message.Payload = payload
	if !session.Send(message) {
		session.Close()
		log.Warning("Dropped session because of broker to client buffer overflow. %v", session.Tag)
	}
}

func randomString() string {
	r := make([]byte, 128/8)
	rand.Read(r)
	return base64.StdEncoding.EncodeToString(r)
}

// sessionGaugeStart starts the gauge for a given session. It returns a new
// function which ends the gauge for the given session. Usually one invokes
// sessionGaugeStart and calls the returned function in a defer statement.
func sessionGaugeStart(session *sockjs.Session) (sessionGaugeEnd func()) {
	changeClientsGauge := lifecycle.CreateClientsGauge()
	changeNewClientsGauge := logger.CreateCounterGauge("newClients", logger.NoUnit, true)
	changeWebsocketClientsGauge := logger.CreateCounterGauge("websocketClients", logger.NoUnit, false)

	changeClientsGauge(1)
	changeNewClientsGauge(1)
	if session.IsWebsocket {
		changeWebsocketClientsGauge(1)
	}

	return func() {
		log.Debug("Client disconnected: %v", session.Tag)
		changeClientsGauge(-1)
		if session.IsWebsocket {
			changeWebsocketClientsGauge(-1)
		}
	}
}

func (b *Broker) sockjsSession(session *sockjs.Session) {
	defer log.RecoverAndLog()

	sessionGaugeEnd := sessionGaugeStart(session)
	defer sessionGaugeEnd()

	socketId := randomString()
	session.Tag = socketId
	log.Debug("Client connected: %v", socketId)

	var controlChannel *amqp.Channel
	var lastPayload string
	resetControlChannel := func() {
		if controlChannel != nil {
			controlChannel.Close()
		}
		var err error
		controlChannel, err = b.PublishConn.Channel()
		if err != nil {
			panic(err)
		}
		go func() {
			defer log.RecoverAndLog()

			for amqpErr := range controlChannel.NotifyClose(make(chan *amqp.Error)) {
				if !(strings.Contains(amqpErr.Error(), "NOT_FOUND") && (strings.Contains(amqpErr.Error(), "koding-social-") || strings.Contains(amqpErr.Error(), "auth-"))) {
					log.Warning("AMQP channel: %v Last publish payload: %v", amqpErr.Error(), lastPayload)
				}

				sendToClient(session, "broker.error", map[string]interface{}{"code": amqpErr.Code, "reason": amqpErr.Reason, "server": amqpErr.Server, "recover": amqpErr.Recover})
			}
		}()
	}
	resetControlChannel()
	defer func() { controlChannel.Close() }()

	subscriptions := make(map[string]bool)
	globalMapMutex.Lock()
	socketSubscriptionsMap[socketId] = &subscriptions
	globalMapMutex.Unlock()

	removeFromRouteMap := func(routingKeyPrefix string) {
		routeSessions := routeMap[routingKeyPrefix]
		for i, routeSession := range routeSessions {
			if routeSession == session {
				routeSessions[i] = routeSessions[len(routeSessions)-1]
				routeSessions = routeSessions[:len(routeSessions)-1]
				break
			}
		}
		if len(routeSessions) == 0 {
			delete(routeMap, routingKeyPrefix)
			return
		}
		routeMap[routingKeyPrefix] = routeSessions
	}

	subscribe := func(routingKeyPrefix string) {
		if subscriptions[routingKeyPrefix] {
			log.Warning("Duplicate subscription to same routing key. %v %v", session.Tag, routingKeyPrefix)
			return
		}
		if len(subscriptions) > 0 && len(subscriptions)%2000 == 0 {
			log.Warning("Client with more than %v subscriptions %v", strconv.Itoa(len(subscriptions)), session.Tag)
		}
		routeMap[routingKeyPrefix] = append(routeMap[routingKeyPrefix], session)
		subscriptions[routingKeyPrefix] = true
	}

	unsubscribe := func(routingKeyPrefix string) {
		removeFromRouteMap(routingKeyPrefix)
		delete(subscriptions, routingKeyPrefix)
	}

	defer func() {
		globalMapMutex.Lock()
		for routingKeyPrefix := range subscriptions {
			removeFromRouteMap(routingKeyPrefix)
		}
		globalMapMutex.Unlock()

		time.AfterFunc(5*time.Minute, func() {
			globalMapMutex.Lock()
			delete(socketSubscriptionsMap, socketId)
			globalMapMutex.Unlock()
		})

		for {
			err := controlChannel.Publish(config.Current.Broker.AuthAllExchange, "broker.clientDisconnected", false, false, amqp.Publishing{Body: []byte(socketId)})
			if err == nil {
				break
			}
			if amqpError, isAmqpError := err.(*amqp.Error); !isAmqpError || amqpError.Code != 504 {
				panic(err)
			}
			resetControlChannel()
		}
	}()

	err := controlChannel.Publish(config.Current.Broker.AuthAllExchange, "broker.clientConnected", false, false, amqp.Publishing{Body: []byte(socketId)})
	if err != nil {
		panic(err)
	}

	sendToClient(session, "broker.connected", socketId)

	for data := range session.ReceiveChan {
		if data == nil || session.Closed {
			break
		}

		defer log.RecoverAndLog()

		message := data.(map[string]interface{})
		log.Debug("Received message: %v", message)

		action := message["action"]
		switch action {
		case "subscribe":
			globalMapMutex.Lock()
			defer globalMapMutex.Unlock()
			for _, routingKeyPrefix := range strings.Split(message["routingKeyPrefix"].(string), " ") {
				subscribe(routingKeyPrefix)
			}
			sendToClient(session, "broker.subscribed", message["routingKeyPrefix"])

		case "resubscribe":
			globalMapMutex.Lock()
			defer globalMapMutex.Unlock()
			oldSubscriptions, found := socketSubscriptionsMap[message["socketId"].(string)]
			if found {
				for routingKeyPrefix := range *oldSubscriptions {
					subscribe(routingKeyPrefix)
				}
			}
			sendToClient(session, "broker.resubscribed", found)

		case "unsubscribe":
			globalMapMutex.Lock()
			defer globalMapMutex.Unlock()
			for _, routingKeyPrefix := range strings.Split(message["routingKeyPrefix"].(string), " ") {
				unsubscribe(routingKeyPrefix)
			}

		case "publish":
			exchange := message["exchange"].(string)
			routingKey := message["routingKey"].(string)
			if !strings.HasPrefix(routingKey, "client.") {
				log.Warning("Invalid routing key: message: %v socketId: %v", message, socketId)
				return
			}
			for {
				lastPayload = ""
				err := controlChannel.Publish(exchange, routingKey, false, false, amqp.Publishing{CorrelationId: socketId, Body: []byte(message["payload"].(string))})
				if err == nil {
					lastPayload = message["payload"].(string)
					break
				}
				if amqpError, isAmqpError := err.(*amqp.Error); !isAmqpError || amqpError.Code != 504 {
					log.Warning("payload: %v routing key: %v exchange: %v err: %v", message["payload"], message["routingKey"], message["exchange"], err)
				}
				time.Sleep(time.Second / 4) // penalty for crashing the AMQP channel
				resetControlChannel()
			}

		case "ping":
			sendToClient(session, "broker.pong", nil)

		default:
			log.Warning("Invalid action. message: %v socketId: %v", message, socketId)

		}
	}
}
<|endoftext|>