<commit_before>/*
Copyright 2017 The Kubernetes 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 dns

const (
	// CoreDNSService is the CoreDNS Service manifest
	CoreDNSService = `
apiVersion: v1
kind: Service
metadata:
  labels:
    k8s-app: kube-dns
    kubernetes.io/cluster-service: "true"
    kubernetes.io/name: "CoreDNS"
  name: kube-dns
  namespace: kube-system
  annotations:
    prometheus.io/port: "9153"
    prometheus.io/scrape: "true"
  # Without this resourceVersion value, an update of the Service between versions will yield:
  #   Service "kube-dns" is invalid: metadata.resourceVersion: Invalid value: "": must be specified for an update
  resourceVersion: "0"
spec:
  clusterIP: {{ .DNSIP }}
  ports:
  - name: dns
    port: 53
    protocol: UDP
    targetPort: 53
  - name: dns-tcp
    port: 53
    protocol: TCP
    targetPort: 53
  - name: metrics
    port: 9153
    protocol: TCP
    targetPort: 9153
  selector:
    k8s-app: kube-dns
`

	// CoreDNSDeployment is the CoreDNS Deployment manifest
	CoreDNSDeployment = `
apiVersion: apps/v1
kind: Deployment
metadata:
  name: {{ .DeploymentName }}
  namespace: kube-system
  labels:
    k8s-app: kube-dns
spec:
  replicas: {{ .Replicas }}
  strategy:
    type: RollingUpdate
    rollingUpdate:
      maxUnavailable: 1
  selector:
    matchLabels:
      k8s-app: kube-dns
  template:
    metadata:
      labels:
        k8s-app: kube-dns
    spec:
      priorityClassName: system-cluster-critical
      serviceAccountName: coredns
      tolerations:
      - key: CriticalAddonsOnly
        operator: Exists
      - key: {{ .OldControlPlaneTaintKey }}
        effect: NoSchedule
      - key: {{ .ControlPlaneTaintKey }}
        effect: NoSchedule
      nodeSelector:
        kubernetes.io/os: linux
      containers:
      - name: coredns
        image: {{ .Image }}
        imagePullPolicy: IfNotPresent
        resources:
          limits:
            memory: 170Mi
          requests:
            cpu: 100m
            memory: 70Mi
        args: [ "-conf", "/etc/coredns/Corefile" ]
        volumeMounts:
        - name: config-volume
          mountPath: /etc/coredns
          readOnly: true
        ports:
        - containerPort: 53
          name: dns
          protocol: UDP
        - containerPort: 53
          name: dns-tcp
          protocol: TCP
        - containerPort: 9153
          name: metrics
          protocol: TCP
        livenessProbe:
          httpGet:
            path: /health
            port: 8080
            scheme: HTTP
          initialDelaySeconds: 60
          timeoutSeconds: 5
          successThreshold: 1
          failureThreshold: 5
        readinessProbe:
          httpGet:
            path: /ready
            port: 8181
            scheme: HTTP
        securityContext:
          allowPrivilegeEscalation: false
          capabilities:
            add:
            - NET_BIND_SERVICE
            drop:
            - all
          readOnlyRootFilesystem: true
      dnsPolicy: Default
      volumes:
        - name: config-volume
          configMap:
            name: coredns
            items:
            - key: Corefile
              path: Corefile
`

	// CoreDNSConfigMap is the CoreDNS ConfigMap manifest
	CoreDNSConfigMap = `
apiVersion: v1
kind: ConfigMap
metadata:
  name: coredns
  namespace: kube-system
data:
  Corefile: |
    .:53 {
        errors
        health {
           lameduck 5s
        }
        ready
        kubernetes {{ .DNSDomain }} in-addr.arpa ip6.arpa {
           pods insecure
           fallthrough in-addr.arpa ip6.arpa
           ttl 30
        }
        prometheus :9153
        forward . /etc/resolv.conf {
           max_concurrent 1000
        }
        cache 30
        loop
        reload
        loadbalance
    }
`
	// CoreDNSClusterRole is the CoreDNS ClusterRole manifest
	CoreDNSClusterRole = `
apiVersion: rbac.authorization.k8s.io/v1
kind: ClusterRole
metadata:
  name: system:coredns
rules:
- apiGroups:
  - ""
  resources:
  - endpoints
  - services
  - pods
  - namespaces
  verbs:
  - list
  - watch
- apiGroups:
  - ""
  resources:
  - nodes
  verbs:
  - get
- apiGroups:
  - discovery.k8s.io
  resources:
  - endpointslices
  verbs:
  - list
  - watch
`
	// CoreDNSClusterRoleBinding is the CoreDNS Clusterrolebinding manifest
	CoreDNSClusterRoleBinding = `
apiVersion: rbac.authorization.k8s.io/v1
kind: ClusterRoleBinding
metadata:
  name: system:coredns
roleRef:
  apiGroup: rbac.authorization.k8s.io
  kind: ClusterRole
  name: system:coredns
subjects:
- kind: ServiceAccount
  name: coredns
  namespace: kube-system
`
	// CoreDNSServiceAccount is the CoreDNS ServiceAccount manifest
	CoreDNSServiceAccount = `
apiVersion: v1
kind: ServiceAccount
metadata:
  name: coredns
  namespace: kube-system
`
)
<commit_msg>kubeadm: add the preferred pod anti-affinity for CoreDNS Deployment<commit_after>/*
Copyright 2017 The Kubernetes 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 dns

const (
	// CoreDNSService is the CoreDNS Service manifest
	CoreDNSService = `
apiVersion: v1
kind: Service
metadata:
  labels:
    k8s-app: kube-dns
    kubernetes.io/cluster-service: "true"
    kubernetes.io/name: "CoreDNS"
  name: kube-dns
  namespace: kube-system
  annotations:
    prometheus.io/port: "9153"
    prometheus.io/scrape: "true"
  # Without this resourceVersion value, an update of the Service between versions will yield:
  #   Service "kube-dns" is invalid: metadata.resourceVersion: Invalid value: "": must be specified for an update
  resourceVersion: "0"
spec:
  clusterIP: {{ .DNSIP }}
  ports:
  - name: dns
    port: 53
    protocol: UDP
    targetPort: 53
  - name: dns-tcp
    port: 53
    protocol: TCP
    targetPort: 53
  - name: metrics
    port: 9153
    protocol: TCP
    targetPort: 9153
  selector:
    k8s-app: kube-dns
`

	// CoreDNSDeployment is the CoreDNS Deployment manifest
	CoreDNSDeployment = `
apiVersion: apps/v1
kind: Deployment
metadata:
  name: {{ .DeploymentName }}
  namespace: kube-system
  labels:
    k8s-app: kube-dns
spec:
  replicas: {{ .Replicas }}
  strategy:
    type: RollingUpdate
    rollingUpdate:
      maxUnavailable: 1
  selector:
    matchLabels:
      k8s-app: kube-dns
  template:
    metadata:
      labels:
        k8s-app: kube-dns
    spec:
      priorityClassName: system-cluster-critical
      serviceAccountName: coredns
      affinity:
        podAntiAffinity:
          preferredDuringSchedulingIgnoredDuringExecution:
          - weight: 100
            podAffinityTerm:
              labelSelector:
                matchExpressions:
                - key: k8s-app
                  operator: In
                  values: ["kube-dns"]
              topologyKey: kubernetes.io/hostname
      tolerations:
      - key: CriticalAddonsOnly
        operator: Exists
      - key: {{ .OldControlPlaneTaintKey }}
        effect: NoSchedule
      - key: {{ .ControlPlaneTaintKey }}
        effect: NoSchedule
      nodeSelector:
        kubernetes.io/os: linux
      containers:
      - name: coredns
        image: {{ .Image }}
        imagePullPolicy: IfNotPresent
        resources:
          limits:
            memory: 170Mi
          requests:
            cpu: 100m
            memory: 70Mi
        args: [ "-conf", "/etc/coredns/Corefile" ]
        volumeMounts:
        - name: config-volume
          mountPath: /etc/coredns
          readOnly: true
        ports:
        - containerPort: 53
          name: dns
          protocol: UDP
        - containerPort: 53
          name: dns-tcp
          protocol: TCP
        - containerPort: 9153
          name: metrics
          protocol: TCP
        livenessProbe:
          httpGet:
            path: /health
            port: 8080
            scheme: HTTP
          initialDelaySeconds: 60
          timeoutSeconds: 5
          successThreshold: 1
          failureThreshold: 5
        readinessProbe:
          httpGet:
            path: /ready
            port: 8181
            scheme: HTTP
        securityContext:
          allowPrivilegeEscalation: false
          capabilities:
            add:
            - NET_BIND_SERVICE
            drop:
            - all
          readOnlyRootFilesystem: true
      dnsPolicy: Default
      volumes:
        - name: config-volume
          configMap:
            name: coredns
            items:
            - key: Corefile
              path: Corefile
`

	// CoreDNSConfigMap is the CoreDNS ConfigMap manifest
	CoreDNSConfigMap = `
apiVersion: v1
kind: ConfigMap
metadata:
  name: coredns
  namespace: kube-system
data:
  Corefile: |
    .:53 {
        errors
        health {
           lameduck 5s
        }
        ready
        kubernetes {{ .DNSDomain }} in-addr.arpa ip6.arpa {
           pods insecure
           fallthrough in-addr.arpa ip6.arpa
           ttl 30
        }
        prometheus :9153
        forward . /etc/resolv.conf {
           max_concurrent 1000
        }
        cache 30
        loop
        reload
        loadbalance
    }
`
	// CoreDNSClusterRole is the CoreDNS ClusterRole manifest
	CoreDNSClusterRole = `
apiVersion: rbac.authorization.k8s.io/v1
kind: ClusterRole
metadata:
  name: system:coredns
rules:
- apiGroups:
  - ""
  resources:
  - endpoints
  - services
  - pods
  - namespaces
  verbs:
  - list
  - watch
- apiGroups:
  - ""
  resources:
  - nodes
  verbs:
  - get
- apiGroups:
  - discovery.k8s.io
  resources:
  - endpointslices
  verbs:
  - list
  - watch
`
	// CoreDNSClusterRoleBinding is the CoreDNS Clusterrolebinding manifest
	CoreDNSClusterRoleBinding = `
apiVersion: rbac.authorization.k8s.io/v1
kind: ClusterRoleBinding
metadata:
  name: system:coredns
roleRef:
  apiGroup: rbac.authorization.k8s.io
  kind: ClusterRole
  name: system:coredns
subjects:
- kind: ServiceAccount
  name: coredns
  namespace: kube-system
`
	// CoreDNSServiceAccount is the CoreDNS ServiceAccount manifest
	CoreDNSServiceAccount = `
apiVersion: v1
kind: ServiceAccount
metadata:
  name: coredns
  namespace: kube-system
`
)
<|endoftext|>


<commit_before>package router

import (
	"bytes"
	"encoding/json"
	"errors"
	"net"
	"net/http"
	"net/url"
	"testing"

	"github.com/micro/go-os/config"
	"github.com/micro/go-os/config/source/memory"
)

func TestRouter(t *testing.T) {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer l.Close()

	routes := []Route{
		{
			Request: Request{
				Method: "GET",
				Host:   l.Addr().String(),
				Path:   "/",
			},
			Response: Response{
				StatusCode: 302,
				Header: map[string]string{
					"Location": "http://example.com",
				},
			},
			Weight: 1.0,
		},
		{
			Request: Request{
				Method: "POST",
				Host:   l.Addr().String(),
				Path:   "/bar",
			},
			Response: Response{
				StatusCode: 301,
				Header: map[string]string{
					"Location": "http://foo.bar.com",
				},
			},
			Weight: 1.0,
		},
		{
			Request: Request{
				Method: "GET",
				Host:   l.Addr().String(),
				Path:   "/foobar",
			},
			ProxyURL: URL{
				Scheme: "http",
				Host:   "www.foo.com",
				Path:   "/",
			},
			Weight: 1.0,
			Type:   "proxy",
		},
	}

	apiConfig := map[string]interface{}{
		"api": map[string]interface{}{
			"routes": routes,
		},
	}

	b, _ := json.Marshal(apiConfig)
	m := memory.NewSource()
	m.Update(b)
	conf := config.NewConfig(config.WithSource(m))
	r := NewRouter(Config(conf))

	wr := r.Handler()
	h := wr(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "not found", 404)
	}))

	go http.Serve(l, h)

	ErrRedirect := errors.New("redirect")

	c := &http.Client{
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return ErrRedirect
		},
	}

	for _, route := range routes {
		var rsp *http.Response
		var err error

		switch route.Request.Method {
		case "GET":
			rsp, err = c.Get("http://" + route.Request.Host + route.Request.Path)
		case "POST":
			rsp, err = c.Post("http://"+route.Request.Host+route.Request.Path, "application/json", bytes.NewBuffer(nil))
		}

		if err != nil {
			urlErr, ok := err.(*url.Error)
			if ok && urlErr.Err == ErrRedirect {
				// skip
			} else {
				t.Fatal(err)
			}
		}

		if route.Type == "proxy" {
			if rsp.StatusCode >= 400 {
				t.Fatalf("Expected healthy response got %d", rsp.StatusCode)
			}
			continue
		}

		if rsp.StatusCode != route.Response.StatusCode {
			t.Fatalf("Expected code %d got %d", route.Response.StatusCode, rsp.StatusCode)
		}

		loc := rsp.Header.Get("Location")
		if loc != route.Response.Header["Location"] {
			t.Fatalf("Expected Location %s got %s", route.Response.Header["Location"], loc)
		}
	}
}
<commit_msg>set router url to google<commit_after>package router

import (
	"bytes"
	"encoding/json"
	"errors"
	"net"
	"net/http"
	"net/url"
	"testing"

	"github.com/micro/go-os/config"
	"github.com/micro/go-os/config/source/memory"
)

func TestRouter(t *testing.T) {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer l.Close()

	routes := []Route{
		{
			Request: Request{
				Method: "GET",
				Host:   l.Addr().String(),
				Path:   "/",
			},
			Response: Response{
				StatusCode: 302,
				Header: map[string]string{
					"Location": "http://example.com",
				},
			},
			Weight: 1.0,
		},
		{
			Request: Request{
				Method: "POST",
				Host:   l.Addr().String(),
				Path:   "/bar",
			},
			Response: Response{
				StatusCode: 301,
				Header: map[string]string{
					"Location": "http://foo.bar.com",
				},
			},
			Weight: 1.0,
		},
		{
			Request: Request{
				Method: "GET",
				Host:   l.Addr().String(),
				Path:   "/foobar",
			},
			ProxyURL: URL{
				Scheme: "https",
				Host:   "www.google.com",
				Path:   "/",
			},
			Weight: 1.0,
			Type:   "proxy",
		},
	}

	apiConfig := map[string]interface{}{
		"api": map[string]interface{}{
			"routes": routes,
		},
	}

	b, _ := json.Marshal(apiConfig)
	m := memory.NewSource()
	m.Update(b)
	conf := config.NewConfig(config.WithSource(m))
	r := NewRouter(Config(conf))

	wr := r.Handler()
	h := wr(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "not found", 404)
	}))

	go http.Serve(l, h)

	ErrRedirect := errors.New("redirect")

	c := &http.Client{
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return ErrRedirect
		},
	}

	for _, route := range routes {
		var rsp *http.Response
		var err error

		switch route.Request.Method {
		case "GET":
			rsp, err = c.Get("http://" + route.Request.Host + route.Request.Path)
		case "POST":
			rsp, err = c.Post("http://"+route.Request.Host+route.Request.Path, "application/json", bytes.NewBuffer(nil))
		}

		if err != nil {
			urlErr, ok := err.(*url.Error)
			if ok && urlErr.Err == ErrRedirect {
				// skip
			} else {
				t.Fatal(err)
			}
		}

		if route.Type == "proxy" {
			if rsp.StatusCode >= 400 {
				t.Fatalf("Expected healthy response got %d", rsp.StatusCode)
			}
			continue
		}

		if rsp.StatusCode != route.Response.StatusCode {
			t.Fatalf("Expected code %d got %d", route.Response.StatusCode, rsp.StatusCode)
		}

		loc := rsp.Header.Get("Location")
		if loc != route.Response.Header["Location"] {
			t.Fatalf("Expected Location %s got %s", route.Response.Header["Location"], loc)
		}
	}
}
<|endoftext|>


<commit_before>package core

import (
	"context"
	"errors"
	"fmt"
	"io"
	"math/rand"
	"sync"
	"time"

	math2 "github.com/ipfs/go-ipfs/thirdparty/math2"
	lgbl "gx/ipfs/QmZ4zF1mBrt8C2mSCM4ZYE4aAnv78f7GvrzufJC4G5tecK/go-libp2p-loggables"

	peer "gx/ipfs/QmQsErDt8Qgw1XrsXf2BpEzDgGWtB1YLsTAARBup5b6B9W/go-libp2p-peer"
	goprocess "gx/ipfs/QmSF8fPo3jgVBAy8fpdjjYqgG87dkJgUprRBHRd2tmfgpP/goprocess"
	procctx "gx/ipfs/QmSF8fPo3jgVBAy8fpdjjYqgG87dkJgUprRBHRd2tmfgpP/goprocess/context"
	periodicproc "gx/ipfs/QmSF8fPo3jgVBAy8fpdjjYqgG87dkJgUprRBHRd2tmfgpP/goprocess/periodic"
	config "gx/ipfs/QmYVqYJTVjetcf1guieEgWpK1PZtHPytP624vKzTF1P3r2/go-ipfs-config"
	inet "gx/ipfs/QmZNJyx9GGCX4GeuHnLB8fxaxMLs4MjTjHokxfQcCd6Nve/go-libp2p-net"
	pstore "gx/ipfs/Qmda4cPRvSRyox3SqgJN6DfSZGU5TtHufPTp9uXjFj71X6/go-libp2p-peerstore"
	host "gx/ipfs/QmeMYW7Nj8jnnEfs9qhm7SxKkoDPUWXu3MsxX6BFwz34tf/go-libp2p-host"
)

// ErrNotEnoughBootstrapPeers signals that we do not have enough bootstrap
// peers to bootstrap correctly.
var ErrNotEnoughBootstrapPeers = errors.New("not enough bootstrap peers to bootstrap")

// BootstrapConfig specifies parameters used in an IpfsNode's network
// bootstrapping process.
type BootstrapConfig struct {

	// MinPeerThreshold governs whether to bootstrap more connections. If the
	// node has less open connections than this number, it will open connections
	// to the bootstrap nodes. From there, the routing system should be able
	// to use the connections to the bootstrap nodes to connect to even more
	// peers. Routing systems like the IpfsDHT do so in their own Bootstrap
	// process, which issues random queries to find more peers.
	MinPeerThreshold int

	// Period governs the periodic interval at which the node will
	// attempt to bootstrap. The bootstrap process is not very expensive, so
	// this threshold can afford to be small (<=30s).
	Period time.Duration

	// ConnectionTimeout determines how long to wait for a bootstrap
	// connection attempt before cancelling it.
	ConnectionTimeout time.Duration

	// BootstrapPeers is a function that returns a set of bootstrap peers
	// for the bootstrap process to use. This makes it possible for clients
	// to control the peers the process uses at any moment.
	BootstrapPeers func() []pstore.PeerInfo
}

// DefaultBootstrapConfig specifies default sane parameters for bootstrapping.
var DefaultBootstrapConfig = BootstrapConfig{
	MinPeerThreshold:  4,
	Period:            30 * time.Second,
	ConnectionTimeout: (30 * time.Second) / 3, // Perod / 3
}

func BootstrapConfigWithPeers(pis []pstore.PeerInfo) BootstrapConfig {
	cfg := DefaultBootstrapConfig
	cfg.BootstrapPeers = func() []pstore.PeerInfo {
		return pis
	}
	return cfg
}

// Bootstrap kicks off IpfsNode bootstrapping. This function will periodically
// check the number of open connections and -- if there are too few -- initiate
// connections to well-known bootstrap peers. It also kicks off subsystem
// bootstrapping (i.e. routing).
func Bootstrap(n *IpfsNode, cfg BootstrapConfig) (io.Closer, error) {

	// make a signal to wait for one bootstrap round to complete.
	doneWithRound := make(chan struct{})

	// the periodic bootstrap function -- the connection supervisor
	periodic := func(worker goprocess.Process) {
		ctx := procctx.OnClosingContext(worker)
		defer log.EventBegin(ctx, "periodicBootstrap", n.Identity).Done()

		if err := bootstrapRound(ctx, n.PeerHost, cfg); err != nil {
			log.Event(ctx, "bootstrapError", n.Identity, lgbl.Error(err))
			log.Debugf("%s bootstrap error: %s", n.Identity, err)
		}

		<-doneWithRound
	}

	// kick off the node's periodic bootstrapping
	proc := periodicproc.Tick(cfg.Period, periodic)
	proc.Go(periodic) // run one right now.

	// kick off Routing.Bootstrap
	if n.Routing != nil {
		ctx := procctx.OnClosingContext(proc)
		if err := n.Routing.Bootstrap(ctx); err != nil {
			proc.Close()
			return nil, err
		}
	}

	doneWithRound <- struct{}{}
	close(doneWithRound) // it no longer blocks periodic
	return proc, nil
}

func bootstrapRound(ctx context.Context, host host.Host, cfg BootstrapConfig) error {

	ctx, cancel := context.WithTimeout(ctx, cfg.ConnectionTimeout)
	defer cancel()
	id := host.ID()

	// get bootstrap peers from config. retrieving them here makes
	// sure we remain observant of changes to client configuration.
	peers := cfg.BootstrapPeers()

	// determine how many bootstrap connections to open
	connected := host.Network().Peers()
	if len(connected) >= cfg.MinPeerThreshold {
		log.Event(ctx, "bootstrapSkip", id)
		log.Debugf("%s core bootstrap skipped -- connected to %d (> %d) nodes",
			id, len(connected), cfg.MinPeerThreshold)
		return nil
	}
	numToDial := cfg.MinPeerThreshold - len(connected)

	// filter out bootstrap nodes we are already connected to
	var notConnected []pstore.PeerInfo
	for _, p := range peers {
		if host.Network().Connectedness(p.ID) != inet.Connected {
			notConnected = append(notConnected, p)
		}
	}

	// if connected to all bootstrap peer candidates, exit
	if len(notConnected) < 1 {
		log.Debugf("%s no more bootstrap peers to create %d connections", id, numToDial)
		return ErrNotEnoughBootstrapPeers
	}

	// connect to a random susbset of bootstrap candidates
	randSubset := randomSubsetOfPeers(notConnected, numToDial)

	defer log.EventBegin(ctx, "bootstrapStart", id).Done()
	log.Debugf("%s bootstrapping to %d nodes: %s", id, numToDial, randSubset)
	return bootstrapConnect(ctx, host, randSubset)
}

func bootstrapConnect(ctx context.Context, ph host.Host, peers []pstore.PeerInfo) error {
	if len(peers) < 1 {
		return ErrNotEnoughBootstrapPeers
	}

	errs := make(chan error, len(peers))
	var wg sync.WaitGroup
	for _, p := range peers {

		// performed asynchronously because when performed synchronously, if
		// one `Connect` call hangs, subsequent calls are more likely to
		// fail/abort due to an expiring context.
		// Also, performed asynchronously for dial speed.

		wg.Add(1)
		go func(p pstore.PeerInfo) {
			defer wg.Done()
			defer log.EventBegin(ctx, "bootstrapDial", ph.ID(), p.ID).Done()
			log.Debugf("%s bootstrapping to %s", ph.ID(), p.ID)

			ph.Peerstore().AddAddrs(p.ID, p.Addrs, pstore.PermanentAddrTTL)
			if err := ph.Connect(ctx, p); err != nil {
				log.Event(ctx, "bootstrapDialFailed", p.ID)
				log.Debugf("failed to bootstrap with %v: %s", p.ID, err)
				errs <- err
				return
			}
			log.Event(ctx, "bootstrapDialSuccess", p.ID)
			log.Infof("bootstrapped with %v", p.ID)
		}(p)
	}
	wg.Wait()

	// our failure condition is when no connection attempt succeeded.
	// So drain the errs channel, counting the results.
	close(errs)
	count := 0
	var err error
	for err = range errs {
		if err != nil {
			count++
		}
	}
	if count == len(peers) {
		return fmt.Errorf("failed to bootstrap. %s", err)
	}
	return nil
}

func toPeerInfos(bpeers []config.BootstrapPeer) []pstore.PeerInfo {
	pinfos := make(map[peer.ID]*pstore.PeerInfo)
	for _, bootstrap := range bpeers {
		pinfo, ok := pinfos[bootstrap.ID()]
		if !ok {
			pinfo = new(pstore.PeerInfo)
			pinfos[bootstrap.ID()] = pinfo
			pinfo.ID = bootstrap.ID()
		}

		pinfo.Addrs = append(pinfo.Addrs, bootstrap.Transport())
	}

	var peers []pstore.PeerInfo
	for _, pinfo := range pinfos {
		peers = append(peers, *pinfo)
	}

	return peers
}

func randomSubsetOfPeers(in []pstore.PeerInfo, max int) []pstore.PeerInfo {
	n := math2.IntMin(max, len(in))
	var out []pstore.PeerInfo
	for _, val := range rand.Perm(len(in)) {
		out = append(out, in[val])
		if len(out) >= n {
			break
		}
	}
	return out
}
<commit_msg>add warning when no bootstrap in config<commit_after>package core

import (
	"context"
	"errors"
	"fmt"
	"io"
	"math/rand"
	"sync"
	"time"

	math2 "github.com/ipfs/go-ipfs/thirdparty/math2"
	lgbl "gx/ipfs/QmZ4zF1mBrt8C2mSCM4ZYE4aAnv78f7GvrzufJC4G5tecK/go-libp2p-loggables"

	peer "gx/ipfs/QmQsErDt8Qgw1XrsXf2BpEzDgGWtB1YLsTAARBup5b6B9W/go-libp2p-peer"
	goprocess "gx/ipfs/QmSF8fPo3jgVBAy8fpdjjYqgG87dkJgUprRBHRd2tmfgpP/goprocess"
	procctx "gx/ipfs/QmSF8fPo3jgVBAy8fpdjjYqgG87dkJgUprRBHRd2tmfgpP/goprocess/context"
	periodicproc "gx/ipfs/QmSF8fPo3jgVBAy8fpdjjYqgG87dkJgUprRBHRd2tmfgpP/goprocess/periodic"
	config "gx/ipfs/QmYVqYJTVjetcf1guieEgWpK1PZtHPytP624vKzTF1P3r2/go-ipfs-config"
	inet "gx/ipfs/QmZNJyx9GGCX4GeuHnLB8fxaxMLs4MjTjHokxfQcCd6Nve/go-libp2p-net"
	pstore "gx/ipfs/Qmda4cPRvSRyox3SqgJN6DfSZGU5TtHufPTp9uXjFj71X6/go-libp2p-peerstore"
	host "gx/ipfs/QmeMYW7Nj8jnnEfs9qhm7SxKkoDPUWXu3MsxX6BFwz34tf/go-libp2p-host"
)

// ErrNotEnoughBootstrapPeers signals that we do not have enough bootstrap
// peers to bootstrap correctly.
var ErrNotEnoughBootstrapPeers = errors.New("not enough bootstrap peers to bootstrap")

// BootstrapConfig specifies parameters used in an IpfsNode's network
// bootstrapping process.
type BootstrapConfig struct {

	// MinPeerThreshold governs whether to bootstrap more connections. If the
	// node has less open connections than this number, it will open connections
	// to the bootstrap nodes. From there, the routing system should be able
	// to use the connections to the bootstrap nodes to connect to even more
	// peers. Routing systems like the IpfsDHT do so in their own Bootstrap
	// process, which issues random queries to find more peers.
	MinPeerThreshold int

	// Period governs the periodic interval at which the node will
	// attempt to bootstrap. The bootstrap process is not very expensive, so
	// this threshold can afford to be small (<=30s).
	Period time.Duration

	// ConnectionTimeout determines how long to wait for a bootstrap
	// connection attempt before cancelling it.
	ConnectionTimeout time.Duration

	// BootstrapPeers is a function that returns a set of bootstrap peers
	// for the bootstrap process to use. This makes it possible for clients
	// to control the peers the process uses at any moment.
	BootstrapPeers func() []pstore.PeerInfo
}

// DefaultBootstrapConfig specifies default sane parameters for bootstrapping.
var DefaultBootstrapConfig = BootstrapConfig{
	MinPeerThreshold:  4,
	Period:            30 * time.Second,
	ConnectionTimeout: (30 * time.Second) / 3, // Perod / 3
}

func BootstrapConfigWithPeers(pis []pstore.PeerInfo) BootstrapConfig {
	cfg := DefaultBootstrapConfig
	cfg.BootstrapPeers = func() []pstore.PeerInfo {
		return pis
	}
	return cfg
}

// Bootstrap kicks off IpfsNode bootstrapping. This function will periodically
// check the number of open connections and -- if there are too few -- initiate
// connections to well-known bootstrap peers. It also kicks off subsystem
// bootstrapping (i.e. routing).
func Bootstrap(n *IpfsNode, cfg BootstrapConfig) (io.Closer, error) {

	// make a signal to wait for one bootstrap round to complete.
	doneWithRound := make(chan struct{})

	// the periodic bootstrap function -- the connection supervisor
	periodic := func(worker goprocess.Process) {
		ctx := procctx.OnClosingContext(worker)
		defer log.EventBegin(ctx, "periodicBootstrap", n.Identity).Done()

		if err := bootstrapRound(ctx, n.PeerHost, cfg); err != nil {
			log.Event(ctx, "bootstrapError", n.Identity, lgbl.Error(err))
			log.Debugf("%s bootstrap error: %s", n.Identity, err)
		}

		<-doneWithRound
	}

	// kick off the node's periodic bootstrapping
	proc := periodicproc.Tick(cfg.Period, periodic)
	proc.Go(periodic) // run one right now.

	// kick off Routing.Bootstrap
	if n.Routing != nil {
		ctx := procctx.OnClosingContext(proc)
		if err := n.Routing.Bootstrap(ctx); err != nil {
			proc.Close()
			return nil, err
		}
	}

	doneWithRound <- struct{}{}
	close(doneWithRound) // it no longer blocks periodic
	return proc, nil
}

func bootstrapRound(ctx context.Context, host host.Host, cfg BootstrapConfig) error {

	ctx, cancel := context.WithTimeout(ctx, cfg.ConnectionTimeout)
	defer cancel()
	id := host.ID()

	// get bootstrap peers from config. retrieving them here makes
	// sure we remain observant of changes to client configuration.
	peers := cfg.BootstrapPeers()
	if len(peers) == 0 {
		log.Warning("no bootstrap in the configuration file")
	}
	// determine how many bootstrap connections to open
	connected := host.Network().Peers()
	if len(connected) >= cfg.MinPeerThreshold {
		log.Event(ctx, "bootstrapSkip", id)
		log.Debugf("%s core bootstrap skipped -- connected to %d (> %d) nodes",
			id, len(connected), cfg.MinPeerThreshold)
		return nil
	}
	numToDial := cfg.MinPeerThreshold - len(connected)

	// filter out bootstrap nodes we are already connected to
	var notConnected []pstore.PeerInfo
	for _, p := range peers {
		if host.Network().Connectedness(p.ID) != inet.Connected {
			notConnected = append(notConnected, p)
		}
	}

	// if connected to all bootstrap peer candidates, exit
	if len(notConnected) < 1 {
		log.Debugf("%s no more bootstrap peers to create %d connections", id, numToDial)
		return ErrNotEnoughBootstrapPeers
	}

	// connect to a random susbset of bootstrap candidates
	randSubset := randomSubsetOfPeers(notConnected, numToDial)

	defer log.EventBegin(ctx, "bootstrapStart", id).Done()
	log.Debugf("%s bootstrapping to %d nodes: %s", id, numToDial, randSubset)
	return bootstrapConnect(ctx, host, randSubset)
}

func bootstrapConnect(ctx context.Context, ph host.Host, peers []pstore.PeerInfo) error {
	if len(peers) < 1 {
		return ErrNotEnoughBootstrapPeers
	}

	errs := make(chan error, len(peers))
	var wg sync.WaitGroup
	for _, p := range peers {

		// performed asynchronously because when performed synchronously, if
		// one `Connect` call hangs, subsequent calls are more likely to
		// fail/abort due to an expiring context.
		// Also, performed asynchronously for dial speed.

		wg.Add(1)
		go func(p pstore.PeerInfo) {
			defer wg.Done()
			defer log.EventBegin(ctx, "bootstrapDial", ph.ID(), p.ID).Done()
			log.Debugf("%s bootstrapping to %s", ph.ID(), p.ID)

			ph.Peerstore().AddAddrs(p.ID, p.Addrs, pstore.PermanentAddrTTL)
			if err := ph.Connect(ctx, p); err != nil {
				log.Event(ctx, "bootstrapDialFailed", p.ID)
				log.Debugf("failed to bootstrap with %v: %s", p.ID, err)
				errs <- err
				return
			}
			log.Event(ctx, "bootstrapDialSuccess", p.ID)
			log.Infof("bootstrapped with %v", p.ID)
		}(p)
	}
	wg.Wait()

	// our failure condition is when no connection attempt succeeded.
	// So drain the errs channel, counting the results.
	close(errs)
	count := 0
	var err error
	for err = range errs {
		if err != nil {
			count++
		}
	}
	if count == len(peers) {
		return fmt.Errorf("failed to bootstrap. %s", err)
	}
	return nil
}

func toPeerInfos(bpeers []config.BootstrapPeer) []pstore.PeerInfo {
	pinfos := make(map[peer.ID]*pstore.PeerInfo)
	for _, bootstrap := range bpeers {
		pinfo, ok := pinfos[bootstrap.ID()]
		if !ok {
			pinfo = new(pstore.PeerInfo)
			pinfos[bootstrap.ID()] = pinfo
			pinfo.ID = bootstrap.ID()
		}

		pinfo.Addrs = append(pinfo.Addrs, bootstrap.Transport())
	}

	var peers []pstore.PeerInfo
	for _, pinfo := range pinfos {
		peers = append(peers, *pinfo)
	}

	return peers
}

func randomSubsetOfPeers(in []pstore.PeerInfo, max int) []pstore.PeerInfo {
	n := math2.IntMin(max, len(in))
	var out []pstore.PeerInfo
	for _, val := range rand.Perm(len(in)) {
		out = append(out, in[val])
		if len(out) >= n {
			break
		}
	}
	return out
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/binary"
	"hash"
	"hash/crc64"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
)

func directoryScanner(directoryScanQueue chan string, fileReadQueue chan string, blockQueue chan block, excludePatterns []string, workInProgress *sync.WaitGroup) {
	for directoryPath := range directoryScanQueue {
		if strings.HasPrefix(directoryPath, "/") {
			logger.Fatalln("unable to create archive with absolute path reference:", directoryPath)
		}
		if verbose {
			logger.Println(directoryPath)
		}

		directory, err := os.Open(directoryPath)
		if err != nil {
			logger.Println("directory read error:", err.Error())
			workInProgress.Done()
			continue
		}

		uid, gid, mode := getModeOwnership(directory)
		blockQueue <- block{directoryPath, 0, nil, blockTypeDirectory, uid, gid, mode}

		for fileName := range readdirnames(directory) {
			filePath := filepath.Join(directoryPath, fileName)

			excludeFile := false
			for _, excludePattern := range excludePatterns {
				match, err := filepath.Match(excludePattern, filePath)
				if err == nil && match {
					excludeFile = true
					break
				}
			}
			if excludeFile {
				logger.Println("skipping excluded file", filePath)
				continue
			}

			fileInfo, err := os.Lstat(filePath)
			if err != nil {
				logger.Println("unable to lstat file", err.Error())
				continue
			} else if (fileInfo.Mode() & os.ModeSymlink) != 0 {
				logger.Println("skipping symbolic link", filePath)
				continue
			}

			workInProgress.Add(1)
			if fileInfo.IsDir() {
				// Sending to directoryScanQueue can block if it's full; since
				// we're also the goroutine responsible for reading from it,
				// this could cause a deadlock.  We break that deadlock by
				// performing the send in a goroutine, where it can block
				// safely.  This does have the side-effect that
				// directoryScanQueue's max size is pretty much ineffective...
				// but that's better than a deadlock.
				go func(filePath string) {
					directoryScanQueue <- filePath
				}(filePath)
			} else {
				fileReadQueue <- filePath
			}
		}

		directory.Close()
		workInProgress.Done()
	}
}

func getModeOwnership(file *os.File) (int, int, os.FileMode) {
	var uid int = 0
	var gid int = 0
	var mode os.FileMode = 0
	fi, err := file.Stat()
	if err != nil {
		logger.Println("file stat error; uid/gid/mode will be incorrect:", err.Error())
	} else {
		mode = fi.Mode()
		stat_t := fi.Sys().(*syscall.Stat_t)
		if stat_t != nil {
			uid = int(stat_t.Uid)
			gid = int(stat_t.Gid)
		} else {
			logger.Println("unable to find file uid/gid")
		}
	}
	return uid, gid, mode
}

func fileReader(fileReadQueue <-chan string, blockQueue chan block, workInProgress *sync.WaitGroup) {
	for filePath := range fileReadQueue {
		if verbose {
			logger.Println(filePath)
		}

		file, err := os.Open(filePath)
		if err == nil {

			uid, gid, mode := getModeOwnership(file)
			blockQueue <- block{filePath, 0, nil, blockTypeStartOfFile, uid, gid, mode}

			bufferedFile := bufio.NewReader(file)

			for {
				buffer := make([]byte, blockSize)
				bytesRead, err := bufferedFile.Read(buffer)
				if err == io.EOF {
					break
				} else if err != nil {
					logger.Println("file read error; file contents will be incomplete:", err.Error())
					break
				}

				blockQueue <- block{filePath, uint16(bytesRead), buffer, blockTypeData, 0, 0, 0}
			}

			blockQueue <- block{filePath, 0, nil, blockTypeEndOfFile, 0, 0, 0}
			file.Close()
		} else {
			logger.Println("file open error:", err.Error())
		}

		workInProgress.Done()
	}
}

func archiveWriter(output io.Writer, blockQueue <-chan block) {
	hash := crc64.New(crc64.MakeTable(crc64.ECMA))
	output = io.MultiWriter(output, hash)
	blockCount := 0
	blockType := make([]byte, 1)

	_, err := output.Write(fastArchiverHeader)
	if err != nil {
		logger.Fatalln("Archive write error:", err.Error())
	}

	for block := range blockQueue {
		filePath := []byte(block.filePath)
		err = binary.Write(output, binary.BigEndian, uint16(len(filePath)))
		if err == nil {
			_, err = output.Write(filePath)
		}
		if err == nil {
			blockType[0] = byte(block.blockType)
			_, err = output.Write(blockType)
		}
		if err == nil {
			switch block.blockType {
			case blockTypeDirectory, blockTypeStartOfFile:
				err = binary.Write(output, binary.BigEndian, uint32(block.uid))
				if err == nil {
					err = binary.Write(output, binary.BigEndian, uint32(block.gid))
				}
				if err == nil {
					err = binary.Write(output, binary.BigEndian, block.mode)
				}
			case blockTypeEndOfFile:
				// Nothing to write aside from the block type
			case blockTypeData:
				err = binary.Write(output, binary.BigEndian, uint16(block.numBytes))
				if err == nil {
					_, err = output.Write(block.buffer[:block.numBytes])
				}
			default:
				logger.Panicln("Unexpected block type")
			}
		}

		blockCount += 1
		if err == nil && (blockCount % 1000) == 0 {
			err = writeChecksumBlock(hash, output, blockType)
		}

		if err != nil {
			logger.Fatalln("Archive write error:", err.Error())
		}
	}

	err = writeChecksumBlock(hash, output, blockType)
	if err != nil {
		logger.Fatalln("Archive write error:", err.Error())
	}
}

func writeChecksumBlock(hash hash.Hash64, output io.Writer, blockType []byte) (error) {
	// file path length... zero
	err := binary.Write(output, binary.BigEndian, uint16(0))
	if err == nil {
		blockType[0] = byte(blockTypeChecksum)
		_, err = output.Write(blockType)
	}
	if err == nil {
		err = binary.Write(output, binary.BigEndian, hash.Sum64())
	}
	return err
}

// Wrapper for Readdirnames that converts it into a generator-style method.
func readdirnames(dir *os.File) chan string {
	retval := make(chan string, 256)
	go func(dir *os.File) {
		for {
			names, err := dir.Readdirnames(256)
			if err == io.EOF {
				break
			} else if err != nil {
				logger.Println("error reading directory:", err.Error())
			}
			for _, name := range names {
				retval <- name
			}
		}
		close(retval)
	}(dir)
	return retval
}
<commit_msg>Separate block writing out of archiveWriter method<commit_after>package main

import (
	"bufio"
	"encoding/binary"
	"hash"
	"hash/crc64"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
)

func directoryScanner(directoryScanQueue chan string, fileReadQueue chan string, blockQueue chan block, excludePatterns []string, workInProgress *sync.WaitGroup) {
	for directoryPath := range directoryScanQueue {
		if strings.HasPrefix(directoryPath, "/") {
			logger.Fatalln("unable to create archive with absolute path reference:", directoryPath)
		}
		if verbose {
			logger.Println(directoryPath)
		}

		directory, err := os.Open(directoryPath)
		if err != nil {
			logger.Println("directory read error:", err.Error())
			workInProgress.Done()
			continue
		}

		uid, gid, mode := getModeOwnership(directory)
		blockQueue <- block{directoryPath, 0, nil, blockTypeDirectory, uid, gid, mode}

		for fileName := range readdirnames(directory) {
			filePath := filepath.Join(directoryPath, fileName)

			excludeFile := false
			for _, excludePattern := range excludePatterns {
				match, err := filepath.Match(excludePattern, filePath)
				if err == nil && match {
					excludeFile = true
					break
				}
			}
			if excludeFile {
				logger.Println("skipping excluded file", filePath)
				continue
			}

			fileInfo, err := os.Lstat(filePath)
			if err != nil {
				logger.Println("unable to lstat file", err.Error())
				continue
			} else if (fileInfo.Mode() & os.ModeSymlink) != 0 {
				logger.Println("skipping symbolic link", filePath)
				continue
			}

			workInProgress.Add(1)
			if fileInfo.IsDir() {
				// Sending to directoryScanQueue can block if it's full; since
				// we're also the goroutine responsible for reading from it,
				// this could cause a deadlock.  We break that deadlock by
				// performing the send in a goroutine, where it can block
				// safely.  This does have the side-effect that
				// directoryScanQueue's max size is pretty much ineffective...
				// but that's better than a deadlock.
				go func(filePath string) {
					directoryScanQueue <- filePath
				}(filePath)
			} else {
				fileReadQueue <- filePath
			}
		}

		directory.Close()
		workInProgress.Done()
	}
}

func getModeOwnership(file *os.File) (int, int, os.FileMode) {
	var uid int = 0
	var gid int = 0
	var mode os.FileMode = 0
	fi, err := file.Stat()
	if err != nil {
		logger.Println("file stat error; uid/gid/mode will be incorrect:", err.Error())
	} else {
		mode = fi.Mode()
		stat_t := fi.Sys().(*syscall.Stat_t)
		if stat_t != nil {
			uid = int(stat_t.Uid)
			gid = int(stat_t.Gid)
		} else {
			logger.Println("unable to find file uid/gid")
		}
	}
	return uid, gid, mode
}

func fileReader(fileReadQueue <-chan string, blockQueue chan block, workInProgress *sync.WaitGroup) {
	for filePath := range fileReadQueue {
		if verbose {
			logger.Println(filePath)
		}

		file, err := os.Open(filePath)
		if err == nil {

			uid, gid, mode := getModeOwnership(file)
			blockQueue <- block{filePath, 0, nil, blockTypeStartOfFile, uid, gid, mode}

			bufferedFile := bufio.NewReader(file)

			for {
				buffer := make([]byte, blockSize)
				bytesRead, err := bufferedFile.Read(buffer)
				if err == io.EOF {
					break
				} else if err != nil {
					logger.Println("file read error; file contents will be incomplete:", err.Error())
					break
				}

				blockQueue <- block{filePath, uint16(bytesRead), buffer, blockTypeData, 0, 0, 0}
			}

			blockQueue <- block{filePath, 0, nil, blockTypeEndOfFile, 0, 0, 0}
			file.Close()
		} else {
			logger.Println("file open error:", err.Error())
		}

		workInProgress.Done()
	}
}

func (b *block) writeBlock(output io.Writer) (error) {
	filePath := []byte(b.filePath)
	err := binary.Write(output, binary.BigEndian, uint16(len(filePath)))
	if err == nil {
		_, err = output.Write(filePath)
	}
	if err == nil {
		blockType := []byte{byte(b.blockType)}
	//blockType := make([]byte, 1)
		//blockType[0] = byte(b.blockType)
		_, err = output.Write(blockType)
	}
	if err == nil {
		switch b.blockType {
		case blockTypeDirectory, blockTypeStartOfFile:
			err = binary.Write(output, binary.BigEndian, uint32(b.uid))
			if err == nil {
				err = binary.Write(output, binary.BigEndian, uint32(b.gid))
			}
			if err == nil {
				err = binary.Write(output, binary.BigEndian, b.mode)
			}
		case blockTypeEndOfFile:
			// Nothing to write aside from the block type
		case blockTypeData:
			err = binary.Write(output, binary.BigEndian, uint16(b.numBytes))
			if err == nil {
				_, err = output.Write(b.buffer[:b.numBytes])
			}
		default:
			logger.Panicln("Unexpected block type")
		}
	}
	return err
}

func archiveWriter(output io.Writer, blockQueue <-chan block) {
	hash := crc64.New(crc64.MakeTable(crc64.ECMA))
	output = io.MultiWriter(output, hash)
	blockCount := 0

	_, err := output.Write(fastArchiverHeader)
	if err != nil {
		logger.Fatalln("Archive write error:", err.Error())
	}

	for block := range blockQueue {
		err = block.writeBlock(output)

		blockCount += 1
		if err == nil && (blockCount % 1000) == 0 {
			err = writeChecksumBlock(hash, output)
		}

		if err != nil {
			logger.Fatalln("Archive write error:", err.Error())
		}
	}

	err = writeChecksumBlock(hash, output)
	if err != nil {
		logger.Fatalln("Archive write error:", err.Error())
	}
}

func writeChecksumBlock(hash hash.Hash64, output io.Writer) (error) {
	// file path length... zero
	err := binary.Write(output, binary.BigEndian, uint16(0))
	if err == nil {
		blockType := []byte{byte(blockTypeChecksum)}
		_, err = output.Write(blockType)
	}
	if err == nil {
		err = binary.Write(output, binary.BigEndian, hash.Sum64())
	}
	return err
}

// Wrapper for Readdirnames that converts it into a generator-style method.
func readdirnames(dir *os.File) chan string {
	retval := make(chan string, 256)
	go func(dir *os.File) {
		for {
			names, err := dir.Readdirnames(256)
			if err == io.EOF {
				break
			} else if err != nil {
				logger.Println("error reading directory:", err.Error())
			}
			for _, name := range names {
				retval <- name
			}
		}
		close(retval)
	}(dir)
	return retval
}
<|endoftext|>


<commit_before>package css // import "github.com/tdewolff/parse/css"

import (
	"bytes"
	"io"
	"strconv"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/tdewolff/buffer"
)

func assertTokens(t *testing.T, s string, tokentypes ...TokenType) {
	stringify := helperStringify(t, s)
	z := NewTokenizer(bytes.NewBufferString(s))
	i := 0
	for {
		tt, _ := z.Next()
		if tt == ErrorToken {
			assert.Equal(t, io.EOF, z.Err(), "error must be EOF in "+stringify)
			assert.Equal(t, len(tokentypes), i, "when error occurred we must be at the end in "+stringify)
			break
		} else if tt == WhitespaceToken {
			continue
		}
		assert.False(t, i >= len(tokentypes), "index must not exceed tokentypes size in "+stringify)
		if i < len(tokentypes) {
			assert.Equal(t, tokentypes[i], tt, "tokentypes must match at index "+strconv.Itoa(i)+" in "+stringify)
		}
		i++
	}
	return
}

func helperStringify(t *testing.T, input string) string {
	s := "["
	z := NewTokenizer(bytes.NewBufferString(input))
	for i := 0; i < 10; i++ {
		tt, text := z.Next()
		if tt == ErrorToken {
			s += tt.String() + "('" + z.Err().Error() + "')"
			break
		} else if tt == WhitespaceToken {
			continue
		} else {
			s += tt.String() + "('" + string(text) + "') "
		}
	}
	return s
}

func assertSplitNumberDimension(t *testing.T, x, e1, e2 string) {
	s1, s2, ok := SplitNumberDimension([]byte(x))
	if !ok && e1 == "" && e2 == "" {
		return
	}
	assert.Equal(t, true, ok, "ok must be true in "+x)
	assert.Equal(t, e1, string(s1), "number part must match in "+x)
	assert.Equal(t, e2, string(s2), "dimension part must match in "+x)
}

func assertSplitDataURI(t *testing.T, x, e1, e2 string, eok bool) {
	s1, s2, ok := SplitDataURI([]byte(x))
	assert.Equal(t, eok, ok, "ok must match in "+x)
	assert.Equal(t, e1, string(s1), "mediatype part must match in "+x)
	assert.Equal(t, e2, string(s2), "data part must match in "+x)
}

////////////////////////////////////////////////////////////////

func TestTokens(t *testing.T) {
	assertTokens(t, " ")
	assertTokens(t, "5.2 .4", NumberToken, NumberToken)
	assertTokens(t, "color: red;", IdentToken, ColonToken, IdentToken, SemicolonToken)
	assertTokens(t, "background: url(\"http://x\");", IdentToken, ColonToken, URLToken, SemicolonToken)
	assertTokens(t, "background: URL(x.png);", IdentToken, ColonToken, URLToken, SemicolonToken)
	assertTokens(t, "color: rgb(4, 0%, 5em);", IdentToken, ColonToken, FunctionToken, NumberToken, CommaToken, PercentageToken, CommaToken, DimensionToken, RightParenthesisToken, SemicolonToken)
	assertTokens(t, "body { \"string\" }", IdentToken, LeftBraceToken, StringToken, RightBraceToken)
	assertTokens(t, "body { \"str\\\"ing\" }", IdentToken, LeftBraceToken, StringToken, RightBraceToken)
	assertTokens(t, ".class { }", DelimToken, IdentToken, LeftBraceToken, RightBraceToken)
	assertTokens(t, "#class { }", HashToken, LeftBraceToken, RightBraceToken)
	assertTokens(t, "#class\\#withhash { }", HashToken, LeftBraceToken, RightBraceToken)
	assertTokens(t, "@media print { }", AtKeywordToken, IdentToken, LeftBraceToken, RightBraceToken)
	assertTokens(t, "/*comment*/", CommentToken)
	assertTokens(t, "/*com* /ment*/", CommentToken)
	assertTokens(t, "~= |= ^= $= *=", IncludeMatchToken, DashMatchToken, PrefixMatchToken, SuffixMatchToken, SubstringMatchToken)
	assertTokens(t, "||", ColumnToken)
	assertTokens(t, "<!-- -->", CDOToken, CDCToken)
	assertTokens(t, "U+1234", UnicodeRangeToken)
	assertTokens(t, "5.2 .4 4e-22", NumberToken, NumberToken, NumberToken)

	// unexpected ending
	assertTokens(t, "ident", IdentToken)
	assertTokens(t, "123.", NumberToken, DelimToken)
	assertTokens(t, "\"string", StringToken)
	assertTokens(t, "123/*comment", NumberToken, CommentToken)
	assertTokens(t, "U+1-", IdentToken, NumberToken, DelimToken)

	// unicode
	assertTokens(t, "fooδbar􀀀", IdentToken)
	assertTokens(t, "foo\\æ\\†", IdentToken)
	//assertTokens(t, "foo\x00bar", IdentToken)
	assertTokens(t, "'foo\u554abar'", StringToken)
	assertTokens(t, "\\000026B", IdentToken)
	assertTokens(t, "\\26 B", IdentToken)

	// hacks
	assertTokens(t, `\-\mo\z\-b\i\nd\in\g:\url(//business\i\nfo.co.uk\/labs\/xbl\/xbl\.xml\#xss);`, IdentToken, ColonToken, URLToken, SemicolonToken)
	assertTokens(t, "width/**/:/**/ 40em;", IdentToken, CommentToken, ColonToken, CommentToken, DimensionToken, SemicolonToken)
	assertTokens(t, ":root *> #quince", ColonToken, IdentToken, DelimToken, DelimToken, HashToken)
	assertTokens(t, "html[xmlns*=\"\"]:root", IdentToken, LeftBracketToken, IdentToken, SubstringMatchToken, StringToken, RightBracketToken, ColonToken, IdentToken)
	assertTokens(t, "body:nth-of-type(1)", IdentToken, ColonToken, FunctionToken, NumberToken, RightParenthesisToken)
	assertTokens(t, "color/*\\**/: blue\\9;", IdentToken, CommentToken, ColonToken, IdentToken, SemicolonToken)
	assertTokens(t, "color: blue !ie;", IdentToken, ColonToken, IdentToken, DelimToken, IdentToken, SemicolonToken)

	// coverage
	assertTokens(t, "  \n\r\n\r\"\\\r\n\\\r\"", StringToken)
	assertTokens(t, "U+?????? U+ABCD?? U+ABC-DEF", UnicodeRangeToken, UnicodeRangeToken, UnicodeRangeToken)
	assertTokens(t, "U+? U+A?", IdentToken, DelimToken, DelimToken, IdentToken, DelimToken, IdentToken, DelimToken)
	assertTokens(t, "-5.23 -moz", NumberToken, IdentToken)
	assertTokens(t, "()", LeftParenthesisToken, RightParenthesisToken)
	assertTokens(t, "url( //url  )", URLToken)
	assertTokens(t, "url( ", URLToken)
	assertTokens(t, "url( //url", URLToken)
	assertTokens(t, "url(\")a", URLToken)
	assertTokens(t, "url(a'\\\n)a", BadURLToken, IdentToken)
	assertTokens(t, "url(\"\n)a", BadURLToken, IdentToken)
	assertTokens(t, "url(a h)a", BadURLToken, IdentToken)
	assertTokens(t, "<!- | @4 ## /2", DelimToken, DelimToken, DelimToken, DelimToken, DelimToken, NumberToken, DelimToken, DelimToken, DelimToken, NumberToken)
	assertTokens(t, "\"s\\\n\"", StringToken)
	assertTokens(t, "\"a\\\"b\"", StringToken)
	assertTokens(t, "\"s\n", BadStringToken)
	//assertTokenError(t, "\\\n", ErrBadEscape)

	assert.Equal(t, "Whitespace", WhitespaceToken.String())
	assert.Equal(t, "Empty", EmptyToken.String())
	assert.Equal(t, "Invalid(100)", TokenType(100).String())
}

func TestTokensSmall(t *testing.T) {
	assertTokens(t, "\"abcd", StringToken)
	assertTokens(t, "/*comment", CommentToken)
	assertTokens(t, "U+A-B", UnicodeRangeToken)
	assertTokens(t, "url((", BadURLToken)
	assertTokens(t, "id\u554a", IdentToken)

	buffer.MinBuf = 5
	assertTokens(t, "ab,d,e", IdentToken, CommaToken, IdentToken, CommaToken, IdentToken)
	assertTokens(t, "ab,cd,e", IdentToken, CommaToken, IdentToken, CommaToken, IdentToken)
}

func TestSplitNumberDimension(t *testing.T) {
	assertSplitNumberDimension(t, "5em", "5", "em")
	assertSplitNumberDimension(t, "+5em", "+5", "em")
	assertSplitNumberDimension(t, "-5.01em", "-5.01", "em")
	assertSplitNumberDimension(t, ".2em", ".2", "em")
	assertSplitNumberDimension(t, ".2e-51em", ".2e-51", "em")
	assertSplitNumberDimension(t, "5%", "5", "%")
	assertSplitNumberDimension(t, "5&%", "", "")
}

func TestSplitDataURI(t *testing.T) {
	assertSplitDataURI(t, "url(www.domain.com)", "", "", false)
	assertSplitDataURI(t, "url(data:,)", "text/plain", "", true)
	assertSplitDataURI(t, "url('data:,')", "text/plain", "", true)
	assertSplitDataURI(t, "url(data:text/xml,)", "text/xml", "", true)
	assertSplitDataURI(t, "url(data:,text)", "text/plain", "text", true)
	assertSplitDataURI(t, "url(data:;base64,dGV4dA==)", "text/plain", "text", true)
	assertSplitDataURI(t, "url(data:image/svg+xml,)", "image/svg+xml", "", true)
}

func TestIsIdent(t *testing.T) {
	assert.True(t, IsIdent([]byte("color")))
	assert.False(t, IsIdent([]byte("4.5")))
}

func TestIsUrlUnquoted(t *testing.T) {
	assert.True(t, IsUrlUnquoted([]byte("http://x")))
	assert.False(t, IsUrlUnquoted([]byte(")")))
}
<commit_msg>Test code change<commit_after>package css // import "github.com/tdewolff/parse/css"

import (
	"bytes"
	"io"
	"strconv"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/tdewolff/buffer"
)

func assertTokens(t *testing.T, s string, tokentypes ...TokenType) {
	stringify := helperStringify(t, s)
	z := NewTokenizer(bytes.NewBufferString(s))
	i := 0
	for {
		tt, _ := z.Next()
		if tt == ErrorToken {
			assert.Equal(t, io.EOF, z.Err(), "error must be EOF in "+stringify)
			assert.Equal(t, len(tokentypes), i, "when error occurred we must be at the end in "+stringify)
			break
		} else if tt == WhitespaceToken {
			continue
		}
		assert.False(t, i >= len(tokentypes), "index must not exceed tokentypes size in "+stringify)
		if i < len(tokentypes) {
			assert.Equal(t, tokentypes[i], tt, "tokentypes must match at index "+strconv.Itoa(i)+" in "+stringify)
		}
		i++
	}
	return
}

func helperStringify(t *testing.T, input string) string {
	s := ""
	z := NewTokenizer(bytes.NewBufferString(input))
	for i := 0; i < 10; i++ {
		tt, text := z.Next()
		if tt == ErrorToken {
			s += tt.String() + "('" + z.Err().Error() + "')"
			break
		} else if tt == WhitespaceToken {
			continue
		} else {
			s += tt.String() + "('" + string(text) + "') "
		}
	}
	return s
}

func assertSplitNumberDimension(t *testing.T, x, e1, e2 string) {
	s1, s2, ok := SplitNumberDimension([]byte(x))
	if !ok && e1 == "" && e2 == "" {
		return
	}
	assert.Equal(t, true, ok, "ok must be true in "+x)
	assert.Equal(t, e1, string(s1), "number part must match in "+x)
	assert.Equal(t, e2, string(s2), "dimension part must match in "+x)
}

func assertSplitDataURI(t *testing.T, x, e1, e2 string, eok bool) {
	s1, s2, ok := SplitDataURI([]byte(x))
	assert.Equal(t, eok, ok, "ok must match in "+x)
	assert.Equal(t, e1, string(s1), "mediatype part must match in "+x)
	assert.Equal(t, e2, string(s2), "data part must match in "+x)
}

////////////////////////////////////////////////////////////////

func TestTokens(t *testing.T) {
	assertTokens(t, " ")
	assertTokens(t, "5.2 .4", NumberToken, NumberToken)
	assertTokens(t, "color: red;", IdentToken, ColonToken, IdentToken, SemicolonToken)
	assertTokens(t, "background: url(\"http://x\");", IdentToken, ColonToken, URLToken, SemicolonToken)
	assertTokens(t, "background: URL(x.png);", IdentToken, ColonToken, URLToken, SemicolonToken)
	assertTokens(t, "color: rgb(4, 0%, 5em);", IdentToken, ColonToken, FunctionToken, NumberToken, CommaToken, PercentageToken, CommaToken, DimensionToken, RightParenthesisToken, SemicolonToken)
	assertTokens(t, "body { \"string\" }", IdentToken, LeftBraceToken, StringToken, RightBraceToken)
	assertTokens(t, "body { \"str\\\"ing\" }", IdentToken, LeftBraceToken, StringToken, RightBraceToken)
	assertTokens(t, ".class { }", DelimToken, IdentToken, LeftBraceToken, RightBraceToken)
	assertTokens(t, "#class { }", HashToken, LeftBraceToken, RightBraceToken)
	assertTokens(t, "#class\\#withhash { }", HashToken, LeftBraceToken, RightBraceToken)
	assertTokens(t, "@media print { }", AtKeywordToken, IdentToken, LeftBraceToken, RightBraceToken)
	assertTokens(t, "/*comment*/", CommentToken)
	assertTokens(t, "/*com* /ment*/", CommentToken)
	assertTokens(t, "~= |= ^= $= *=", IncludeMatchToken, DashMatchToken, PrefixMatchToken, SuffixMatchToken, SubstringMatchToken)
	assertTokens(t, "||", ColumnToken)
	assertTokens(t, "<!-- -->", CDOToken, CDCToken)
	assertTokens(t, "U+1234", UnicodeRangeToken)
	assertTokens(t, "5.2 .4 4e-22", NumberToken, NumberToken, NumberToken)

	// unexpected ending
	assertTokens(t, "ident", IdentToken)
	assertTokens(t, "123.", NumberToken, DelimToken)
	assertTokens(t, "\"string", StringToken)
	assertTokens(t, "123/*comment", NumberToken, CommentToken)
	assertTokens(t, "U+1-", IdentToken, NumberToken, DelimToken)

	// unicode
	assertTokens(t, "fooδbar􀀀", IdentToken)
	assertTokens(t, "foo\\æ\\†", IdentToken)
	//assertTokens(t, "foo\x00bar", IdentToken)
	assertTokens(t, "'foo\u554abar'", StringToken)
	assertTokens(t, "\\000026B", IdentToken)
	assertTokens(t, "\\26 B", IdentToken)

	// hacks
	assertTokens(t, `\-\mo\z\-b\i\nd\in\g:\url(//business\i\nfo.co.uk\/labs\/xbl\/xbl\.xml\#xss);`, IdentToken, ColonToken, URLToken, SemicolonToken)
	assertTokens(t, "width/**/:/**/ 40em;", IdentToken, CommentToken, ColonToken, CommentToken, DimensionToken, SemicolonToken)
	assertTokens(t, ":root *> #quince", ColonToken, IdentToken, DelimToken, DelimToken, HashToken)
	assertTokens(t, "html[xmlns*=\"\"]:root", IdentToken, LeftBracketToken, IdentToken, SubstringMatchToken, StringToken, RightBracketToken, ColonToken, IdentToken)
	assertTokens(t, "body:nth-of-type(1)", IdentToken, ColonToken, FunctionToken, NumberToken, RightParenthesisToken)
	assertTokens(t, "color/*\\**/: blue\\9;", IdentToken, CommentToken, ColonToken, IdentToken, SemicolonToken)
	assertTokens(t, "color: blue !ie;", IdentToken, ColonToken, IdentToken, DelimToken, IdentToken, SemicolonToken)

	// coverage
	assertTokens(t, "  \n\r\n\r\"\\\r\n\\\r\"", StringToken)
	assertTokens(t, "U+?????? U+ABCD?? U+ABC-DEF", UnicodeRangeToken, UnicodeRangeToken, UnicodeRangeToken)
	assertTokens(t, "U+? U+A?", IdentToken, DelimToken, DelimToken, IdentToken, DelimToken, IdentToken, DelimToken)
	assertTokens(t, "-5.23 -moz", NumberToken, IdentToken)
	assertTokens(t, "()", LeftParenthesisToken, RightParenthesisToken)
	assertTokens(t, "url( //url  )", URLToken)
	assertTokens(t, "url( ", URLToken)
	assertTokens(t, "url( //url", URLToken)
	assertTokens(t, "url(\")a", URLToken)
	assertTokens(t, "url(a'\\\n)a", BadURLToken, IdentToken)
	assertTokens(t, "url(\"\n)a", BadURLToken, IdentToken)
	assertTokens(t, "url(a h)a", BadURLToken, IdentToken)
	assertTokens(t, "<!- | @4 ## /2", DelimToken, DelimToken, DelimToken, DelimToken, DelimToken, NumberToken, DelimToken, DelimToken, DelimToken, NumberToken)
	assertTokens(t, "\"s\\\n\"", StringToken)
	assertTokens(t, "\"a\\\"b\"", StringToken)
	assertTokens(t, "\"s\n", BadStringToken)
	//assertTokenError(t, "\\\n", ErrBadEscape)

	assert.Equal(t, "Whitespace", WhitespaceToken.String())
	assert.Equal(t, "Empty", EmptyToken.String())
	assert.Equal(t, "Invalid(100)", TokenType(100).String())
}

func TestTokensSmall(t *testing.T) {
	assertTokens(t, "\"abcd", StringToken)
	assertTokens(t, "/*comment", CommentToken)
	assertTokens(t, "U+A-B", UnicodeRangeToken)
	assertTokens(t, "url((", BadURLToken)
	assertTokens(t, "id\u554a", IdentToken)

	buffer.MinBuf = 5
	assertTokens(t, "ab,d,e", IdentToken, CommaToken, IdentToken, CommaToken, IdentToken)
	assertTokens(t, "ab,cd,e", IdentToken, CommaToken, IdentToken, CommaToken, IdentToken)
}

func TestSplitNumberDimension(t *testing.T) {
	assertSplitNumberDimension(t, "5em", "5", "em")
	assertSplitNumberDimension(t, "+5em", "+5", "em")
	assertSplitNumberDimension(t, "-5.01em", "-5.01", "em")
	assertSplitNumberDimension(t, ".2em", ".2", "em")
	assertSplitNumberDimension(t, ".2e-51em", ".2e-51", "em")
	assertSplitNumberDimension(t, "5%", "5", "%")
	assertSplitNumberDimension(t, "5&%", "", "")
}

func TestSplitDataURI(t *testing.T) {
	assertSplitDataURI(t, "url(www.domain.com)", "", "", false)
	assertSplitDataURI(t, "url(data:,)", "text/plain", "", true)
	assertSplitDataURI(t, "url('data:,')", "text/plain", "", true)
	assertSplitDataURI(t, "url(data:text/xml,)", "text/xml", "", true)
	assertSplitDataURI(t, "url(data:,text)", "text/plain", "text", true)
	assertSplitDataURI(t, "url(data:;base64,dGV4dA==)", "text/plain", "text", true)
	assertSplitDataURI(t, "url(data:image/svg+xml,)", "image/svg+xml", "", true)
}

func TestIsIdent(t *testing.T) {
	assert.True(t, IsIdent([]byte("color")))
	assert.False(t, IsIdent([]byte("4.5")))
}

func TestIsUrlUnquoted(t *testing.T) {
	assert.True(t, IsUrlUnquoted([]byte("http://x")))
	assert.False(t, IsUrlUnquoted([]byte(")")))
}
<|endoftext|>


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

package cmd

import (
	"bufio"
	"fmt"
	"math"
	"os"
	"runtime"
	"strconv"
	"strings"

	"github.com/spf13/cobra"
	"gonum.org/v1/gonum/stat"
)

// corrCmd represents the corr command
var corrCmd = &cobra.Command{
	Use:   "corr",
	Short: "calculate Pearson correlation between two columns",
	Long:  "calculate Pearson correlation between two columns",

	Run: func(cmd *cobra.Command, args []string) {
		config := getConfigs(cmd)
		files := getFileList(args)
		runtime.GOMAXPROCS(config.NumCPUs)

		printField := getFlagString(cmd, "fields")
		printIgnore := getFlagBool(cmd, "ignore_nan")
		_ = printIgnore
		printPass := getFlagBool(cmd, "pass")
		printLog := getFlagBool(cmd, "log")
		outFile := config.OutFile

		if config.Tabs {
			config.OutDelimiter = rune('\t')
		}

		outw := os.Stdout
		if outFile != "-" {
			tw, err := os.Create(outFile)
			checkError(err)
			outw = tw
		}
		outfh := bufio.NewWriter(outw)

		defer outfh.Flush()
		defer outw.Close()

		transform := func(x float64) float64 { return x }
		if printLog {
			transform = func(x float64) float64 {
				return math.Log10(x + 1)
			}
		}

		field2col := make(map[string]int)
		Data := make(map[int][]float64)

		targetCols := make(map[int]string)

		for x, tok := range strings.Split(printField, ",") {
			tok = strings.TrimSpace(tok)
			var col int
			if config.NoHeaderRow {
				if len(tok) == 0 {
					fmt.Fprintf(os.Stderr, "No field specified!")
					os.Exit(1)
				}
				pcol, err := strconv.Atoi(tok)
				if err != nil {
					fmt.Fprintf(os.Stderr, "Illegal field number: %s\n", tok)
					os.Exit(1)
				}
				col = pcol - 1
				if col < 0 {
					fmt.Fprintf(os.Stderr, "Illegal field number: %d!\n", pcol)
					os.Exit(1)
				}
				targetCols[col] = tok
			}
			if len(tok) == 0 {
				fmt.Fprintf(os.Stderr, "Empty field specified!\n")
				os.Exit(1)
			}
			targetCols[-(x + 1)] = tok

		}

		for _, file := range files[:1] {
			csvReader, err := newCSVReaderByConfig(config, file)
			checkError(err)
			csvReader.Run()

			isHeaderLine := !config.NoHeaderRow
			for chunk := range csvReader.Ch {
				checkError(chunk.Err)

				for _, record := range chunk.Data {
					if isHeaderLine {
						for i, column := range record {
							field2col[column] = i
						}

						isHeaderLine = false
						if printPass {
							outfh.Write([]byte(strings.Join(record, string(config.OutDelimiter)) + "\n"))
						}
						continue
					} // header

					for col, field := range targetCols {
						i := col
						if !config.NoHeaderRow && i < 0 {
							var ok bool
							i, ok = field2col[field]
							if !ok {
								fmt.Fprintf(os.Stderr, "Invalid field specified: %s\n", field)
								os.Exit(1)
							}
						}
						if printPass {
							outfh.Write([]byte(strings.Join(record, string(config.OutDelimiter)) + "\n"))
						}
						p, err := strconv.ParseFloat(record[i], 64)
						if err == nil {
							Data[i] = append(Data[i], transform(p))
						} else {
							Data[i] = append(Data[i], math.NaN())
						}
					} // field

				} // record
			} //chunk

		} //file

		// Calculate and print correlations:
		seen := make(map[int]map[int]bool)
		for col1, field1 := range targetCols {
			if col1 < 0 {
				col1 = field2col[field1]
			}
			if seen[col1] == nil {
				seen[col1] = make(map[int]bool)
			}
			for col2, field2 := range targetCols {
				if col2 < 0 {
					col2 = field2col[field2]
				}
				if col1 == col2 {
					continue
				}
				if seen[col1][col2] {
					continue
				}
				d1, d2 := Data[col1], Data[col2]
				if printIgnore {
					d1, d2 = removeNaNs(d1, d2)
				}

				pearsonr := stat.Correlation(d1, d2, nil)
				fmt.Fprintf(os.Stderr, "%s%s%s%s%.4f\n", field1, string(config.OutDelimiter), field2, string(config.OutDelimiter), pearsonr)

				seen[col1][col2] = true
				if seen[col2] == nil {
					seen[col2] = make(map[int]bool)
				}
				seen[col2][col1] = true
			} // col2
		} // col1
	},
}

// removeNaNs removes entries from a pair of slices if any of the two values is NaN.
func removeNaNs(d1, d2 []float64) ([]float64, []float64) {
	r1 := make([]float64, 0, len(d1))
	r2 := make([]float64, 0, len(d1))

	for i, x1 := range d1 {
		x2 := d2[i]
		if !math.IsNaN(x1) && !math.IsNaN(x2) {
			r1 = append(r1, x1)
			r2 = append(r2, x2)
		}
	}
	return r1, r2
}

func init() {
	RootCmd.AddCommand(corrCmd)

	corrCmd.Flags().StringP("fields", "f", "", "comma separated fields")
	corrCmd.Flags().BoolP("ignore_nan", "i", false, "Ignore non-numeric fields to avoid returning NaN")
	corrCmd.Flags().BoolP("log", "L", false, "Calcute correlations on Log10 transformed data")
	corrCmd.Flags().BoolP("pass", "x", false, "passthrough mode (forward input to output)")
}
<commit_msg>corr: use all columns if -f is empty.<commit_after>// Copyright © 2019 Oxford Nanopore Technologies.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cmd

import (
	"bufio"
	"fmt"
	"math"
	"os"
	"runtime"
	"strconv"
	"strings"

	"github.com/spf13/cobra"
	"gonum.org/v1/gonum/stat"
)

// corrCmd represents the corr command
var corrCmd = &cobra.Command{
	Use:   "corr",
	Short: "calculate Pearson correlation between two columns",
	Long:  "calculate Pearson correlation between two columns",

	Run: func(cmd *cobra.Command, args []string) {
		config := getConfigs(cmd)
		files := getFileList(args)
		runtime.GOMAXPROCS(config.NumCPUs)

		printField := getFlagString(cmd, "fields")
		printIgnore := getFlagBool(cmd, "ignore_nan")
		_ = printIgnore
		printPass := getFlagBool(cmd, "pass")
		printLog := getFlagBool(cmd, "log")
		outFile := config.OutFile

		if config.Tabs {
			config.OutDelimiter = rune('\t')
		}

		outw := os.Stdout
		if outFile != "-" {
			tw, err := os.Create(outFile)
			checkError(err)
			outw = tw
		}
		outfh := bufio.NewWriter(outw)

		defer outfh.Flush()
		defer outw.Close()

		transform := func(x float64) float64 { return x }
		if printLog {
			transform = func(x float64) float64 {
				return math.Log10(x + 1)
			}
		}

		field2col := make(map[string]int)
		Data := make(map[int][]float64)

		targetCols := make(map[int]string)

		for x, tok := range strings.Split(printField, ",") {
			tok = strings.TrimSpace(tok)
			var col int
			if config.NoHeaderRow {
				if len(tok) == 0 {
					continue
				}
				pcol, err := strconv.Atoi(tok)
				if err != nil {
					fmt.Fprintf(os.Stderr, "Illegal field number: %s\n", tok)
					os.Exit(1)
				}
				col = pcol - 1
				if col < 0 {
					fmt.Fprintf(os.Stderr, "Illegal field number: %d!\n", pcol)
					os.Exit(1)
				}
				targetCols[col] = tok
			}
			if len(tok) != 0 {
				targetCols[-(x + 1)] = tok
			}

		}

		for _, file := range files[:1] {
			csvReader, err := newCSVReaderByConfig(config, file)
			checkError(err)
			csvReader.Run()

			isHeaderLine := !config.NoHeaderRow
			for chunk := range csvReader.Ch {
				checkError(chunk.Err)

				for _, record := range chunk.Data {
					if isHeaderLine {
						for i, column := range record {
							field2col[column] = i
							if printField == "" {
								targetCols[i] = column
							}
						}

						isHeaderLine = false
						if printPass {
							outfh.Write([]byte(strings.Join(record, string(config.OutDelimiter)) + "\n"))
						}
						continue
					} else {
						if len(targetCols) == 0 {
							for i, _ := range record {
								targetCols[i] = strconv.Itoa(i + 1)
							}
						}
					}

					for col, field := range targetCols {
						i := col
						if !config.NoHeaderRow && i < 0 {
							var ok bool
							i, ok = field2col[field]
							if !ok {
								fmt.Fprintf(os.Stderr, "Invalid field specified: %s\n", field)
								os.Exit(1)
							}
						}
						if printPass {
							outfh.Write([]byte(strings.Join(record, string(config.OutDelimiter)) + "\n"))
						}
						p, err := strconv.ParseFloat(record[i], 64)
						if err == nil {
							Data[i] = append(Data[i], transform(p))
						} else {
							Data[i] = append(Data[i], math.NaN())
						}
					} // field

				} // record
			} //chunk

		} //file

		// Calculate and print correlations:
		seen := make(map[int]map[int]bool)
		for col1, field1 := range targetCols {
			if col1 < 0 {
				col1 = field2col[field1]
			}
			if seen[col1] == nil {
				seen[col1] = make(map[int]bool)
			}
			for col2, field2 := range targetCols {
				if col2 < 0 {
					col2 = field2col[field2]
				}
				if col1 == col2 {
					continue
				}
				if seen[col1][col2] {
					continue
				}
				d1, d2 := Data[col1], Data[col2]
				if printIgnore {
					d1, d2 = removeNaNs(d1, d2)
				}

				pearsonr := stat.Correlation(d1, d2, nil)
				fmt.Fprintf(os.Stderr, "%s%s%s%s%.4f\n", field1, string(config.OutDelimiter), field2, string(config.OutDelimiter), pearsonr)

				seen[col1][col2] = true
				if seen[col2] == nil {
					seen[col2] = make(map[int]bool)
				}
				seen[col2][col1] = true
			} // col2
		} // col1
	},
}

// removeNaNs removes entries from a pair of slices if any of the two values is NaN.
func removeNaNs(d1, d2 []float64) ([]float64, []float64) {
	r1 := make([]float64, 0, len(d1))
	r2 := make([]float64, 0, len(d1))

	for i, x1 := range d1 {
		x2 := d2[i]
		if !math.IsNaN(x1) && !math.IsNaN(x2) {
			r1 = append(r1, x1)
			r2 = append(r2, x2)
		}
	}
	return r1, r2
}

func init() {
	RootCmd.AddCommand(corrCmd)

	corrCmd.Flags().StringP("fields", "f", "", "comma separated fields")
	corrCmd.Flags().BoolP("ignore_nan", "i", false, "Ignore non-numeric fields to avoid returning NaN")
	corrCmd.Flags().BoolP("log", "L", false, "Calcute correlations on Log10 transformed data")
	corrCmd.Flags().BoolP("pass", "x", false, "passthrough mode (forward input to output)")
}
<|endoftext|>


<commit_before>package util

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net"
	"os/user"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/skia-dev/glog"
	"go.skia.org/infra/go/util"
	"golang.org/x/crypto/ssh"
)

const (
	KEY_FILE           = "id_rsa"
	WORKER_NUM_KEYWORD = "{{worker_num}}"
)

type workerResp struct {
	hostname string
	output   string
}

func executeCmd(cmd, hostname string, config *ssh.ClientConfig, timeout time.Duration) (string, error) {
	// Dial up TCP connection to remote machine.
	conn, err := net.Dial("tcp", hostname+":22")
	if err != nil {
		return "", fmt.Errorf("Failed to ssh connect to %s. Make sure \"PubkeyAuthentication yes\" is in your sshd_config: %s", hostname, err)
	}
	defer util.Close(conn)
	util.LogErr(conn.SetDeadline(time.Now().Add(timeout)))

	// Create new SSH client connection.
	sshConn, sshChan, req, err := ssh.NewClientConn(conn, hostname+":22", config)
	if err != nil {
		return "", fmt.Errorf("Failed to ssh connect to %s: %s", hostname, err)
	}
	// Use client connection to create new client.
	client := ssh.NewClient(sshConn, sshChan, req)

	// Client connections can support multiple interactive sessions.
	session, err := client.NewSession()
	if err != nil {
		return "", fmt.Errorf("Failed to ssh connect to %s: %s", hostname, err)
	}

	var stdoutBuf bytes.Buffer
	session.Stdout = &stdoutBuf
	if err := session.Run(cmd); err != nil {
		return "", fmt.Errorf("Errored or Timeout out while running \"%s\" on %s: %s", cmd, hostname, err)
	}
	return stdoutBuf.String(), nil
}

func getKeyFile() (key ssh.Signer, err error) {
	usr, _ := user.Current()
	file := usr.HomeDir + "/.ssh/" + KEY_FILE
	buf, err := ioutil.ReadFile(file)
	if err != nil {
		return
	}
	key, err = ssh.ParsePrivateKey(buf)
	if err != nil {
		return
	}
	return
}

// SSH connects to the specified workers and runs the specified command. If the
// command does not complete in the given duration then all remaining workers are
// considered timed out. SSH also automatically substitutes the sequential number
// of the worker for the WORKER_NUM_KEYWORD since it is a common use case.
func SSH(cmd string, workers []string, timeout time.Duration) (map[string]string, error) {
	glog.Infof("Running \"%s\" on %s with timeout of %s", cmd, workers, timeout)
	numWorkers := len(workers)

	// Ensure that the key file exists.
	key, err := getKeyFile()
	if err != nil {
		return nil, fmt.Errorf("Failed to get key file: %s", err)
	}

	// Initialize the structure with the configuration for ssh.
	config := &ssh.ClientConfig{
		User: CtUser,
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(key),
		},
	}

	var wg sync.WaitGroup
	// m protects workersWithOutputs and remainingWorkers
	var m sync.Mutex
	// Will be populated and returned by this function.
	workersWithOutputs := map[string]string{}
	// Keeps track of which workers are still pending.
	remainingWorkers := map[string]int{}

	// Kick off a goroutine on all workers.
	for i, hostname := range workers {
		wg.Add(1)
		m.Lock()
		remainingWorkers[hostname] = 1
		m.Unlock()
		go func(index int, hostname string) {
			defer wg.Done()
			updatedCmd := strings.Replace(cmd, WORKER_NUM_KEYWORD, strconv.Itoa(index+1), -1)
			output, err := executeCmd(updatedCmd, hostname, config, timeout)
			if err != nil {
				glog.Errorf("Could not execute ssh cmd: %s", err)
			}
			m.Lock()
			defer m.Unlock()
			workersWithOutputs[hostname] = output
			delete(remainingWorkers, hostname)
			glog.Infoln()
			glog.Infof("[%d/%d] Worker %s has completed execution", numWorkers-len(remainingWorkers), numWorkers, hostname)
			glog.Infof("Remaining workers: %v", remainingWorkers)
		}(i, hostname)
	}

	wg.Wait()
	glog.Infoln()
	glog.Infof("Finished running \"%s\" on all %d workers", cmd, numWorkers)
	glog.Info("========================================")

	m.Lock()
	defer m.Unlock()
	return workersWithOutputs, nil
}

// RebootWorkers reboots all CT workers and waits for few mins before returning.
func RebootWorkers() {
	if _, err := SSH("sudo reboot", Slaves, REBOOT_TIMEOUT); err != nil {
		glog.Errorf("Got error while rebooting workers: %v", err)
	}
	waitTime := 15 * time.Minute
	glog.Infof("Waiting for %s till all workers come back from reboot", waitTime)
	time.Sleep(waitTime)
}

// RebootAndroidDevices reboots the Android device on all CT workers and waits
// for few mins before returning.
func RebootAndroidDevices() {
	if _, err := SSH("adb reboot", Slaves, REBOOT_TIMEOUT); err != nil {
		glog.Errorf("Got error while rebooting devices: %v", err)
	}
	waitTime := 5 * time.Minute
	glog.Infof("Waiting for %s till all Android devices come back from reboot", waitTime)
	time.Sleep(waitTime)
}
<commit_msg>[CT] Use better algorithm for waiting for reboots<commit_after>package util

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net"
	"os/user"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/skia-dev/glog"
	"go.skia.org/infra/go/util"
	"golang.org/x/crypto/ssh"
)

const (
	KEY_FILE           = "id_rsa"
	WORKER_NUM_KEYWORD = "{{worker_num}}"
)

type workerResp struct {
	hostname string
	output   string
}

func executeCmd(cmd, hostname string, config *ssh.ClientConfig, timeout time.Duration) (string, error) {
	// Dial up TCP connection to remote machine.
	conn, err := net.Dial("tcp", hostname+":22")
	if err != nil {
		return "", fmt.Errorf("Failed to ssh connect to %s. Make sure \"PubkeyAuthentication yes\" is in your sshd_config: %s", hostname, err)
	}
	defer util.Close(conn)
	util.LogErr(conn.SetDeadline(time.Now().Add(timeout)))

	// Create new SSH client connection.
	sshConn, sshChan, req, err := ssh.NewClientConn(conn, hostname+":22", config)
	if err != nil {
		return "", fmt.Errorf("Failed to ssh connect to %s: %s", hostname, err)
	}
	// Use client connection to create new client.
	client := ssh.NewClient(sshConn, sshChan, req)

	// Client connections can support multiple interactive sessions.
	session, err := client.NewSession()
	if err != nil {
		return "", fmt.Errorf("Failed to ssh connect to %s: %s", hostname, err)
	}

	var stdoutBuf bytes.Buffer
	session.Stdout = &stdoutBuf
	if err := session.Run(cmd); err != nil {
		return "", fmt.Errorf("Errored or Timeout out while running \"%s\" on %s: %s", cmd, hostname, err)
	}
	return stdoutBuf.String(), nil
}

func getKeyFile() (key ssh.Signer, err error) {
	usr, _ := user.Current()
	file := usr.HomeDir + "/.ssh/" + KEY_FILE
	buf, err := ioutil.ReadFile(file)
	if err != nil {
		return
	}
	key, err = ssh.ParsePrivateKey(buf)
	if err != nil {
		return
	}
	return
}

// SSH connects to the specified workers and runs the specified command. If the
// command does not complete in the given duration then all remaining workers are
// considered timed out. SSH also automatically substitutes the sequential number
// of the worker for the WORKER_NUM_KEYWORD since it is a common use case.
func SSH(cmd string, workers []string, timeout time.Duration) (map[string]string, error) {
	glog.Infof("Running \"%s\" on %s with timeout of %s", cmd, workers, timeout)
	numWorkers := len(workers)

	// Ensure that the key file exists.
	key, err := getKeyFile()
	if err != nil {
		return nil, fmt.Errorf("Failed to get key file: %s", err)
	}

	// Initialize the structure with the configuration for ssh.
	config := &ssh.ClientConfig{
		User: CtUser,
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(key),
		},
	}

	var wg sync.WaitGroup
	// m protects workersWithOutputs and remainingWorkers
	var m sync.Mutex
	// Will be populated and returned by this function.
	workersWithOutputs := map[string]string{}
	// Keeps track of which workers are still pending.
	remainingWorkers := map[string]int{}

	// Kick off a goroutine on all workers.
	for i, hostname := range workers {
		wg.Add(1)
		m.Lock()
		remainingWorkers[hostname] = 1
		m.Unlock()
		go func(index int, hostname string) {
			defer wg.Done()
			updatedCmd := strings.Replace(cmd, WORKER_NUM_KEYWORD, strconv.Itoa(index+1), -1)
			output, err := executeCmd(updatedCmd, hostname, config, timeout)
			if err != nil {
				glog.Errorf("Could not execute ssh cmd: %s", err)
			}
			m.Lock()
			defer m.Unlock()
			workersWithOutputs[hostname] = output
			delete(remainingWorkers, hostname)
			glog.Infoln()
			glog.Infof("[%d/%d] Worker %s has completed execution", numWorkers-len(remainingWorkers), numWorkers, hostname)
			glog.Infof("Remaining workers: %v", remainingWorkers)
		}(i, hostname)
	}

	wg.Wait()
	glog.Infoln()
	glog.Infof("Finished running \"%s\" on all %d workers", cmd, numWorkers)
	glog.Info("========================================")

	m.Lock()
	defer m.Unlock()
	return workersWithOutputs, nil
}

// RebootWorkers reboots all CT workers and waits for them to return.
func RebootWorkers() {
	if _, err := SSH("sudo reboot", Slaves, REBOOT_TIMEOUT); err != nil {
		glog.Errorf("Got error while rebooting workers: %v", err)
		return
	}
	waitTime := 5 * time.Minute
	glog.Infof("Waiting for %s till all workers come back from reboot", waitTime)

	// Check every 2 mins and timeout after 10 mins.
	ticker := time.NewTicker(2 * time.Minute)
	deadlineTicker := time.NewTicker(10 * time.Minute)
	defer ticker.Stop()
	defer deadlineTicker.Stop()
	for {
		select {
		case <-ticker.C:
			output, err := SSH("uptime", Slaves, REBOOT_TIMEOUT)
			if err != nil {
				glog.Errorf("Got error while checking workers: %v", err)
				return
			}
			if len(output) == NUM_WORKERS_PROD {
				glog.Infof("All workers are back.")
				return
			} else {
				glog.Infof("Got replies from %d/%d slaves. Continuing to wait.", len(output), NUM_WORKERS_PROD)
			}
		case <-deadlineTicker.C:
			fmt.Println("Deadline surpassed so we are done waiting for slaves.")
			return
		}
	}
}

// RebootAndroidDevices reboots the Android device on all CT workers and waits
// for few mins before returning.
func RebootAndroidDevices() {
	if _, err := SSH("adb reboot", Slaves, REBOOT_TIMEOUT); err != nil {
		glog.Errorf("Got error while rebooting devices: %v", err)
	}
	waitTime := 5 * time.Minute
	glog.Infof("Waiting for %s till all Android devices come back from reboot", waitTime)
	time.Sleep(waitTime)
}
<|endoftext|>


<commit_before>package cwriter

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"

	"golang.org/x/crypto/ssh/terminal"
)

var NotATTY = errors.New("not a terminal")

var cuuAndEd = fmt.Sprintf("%c[%%dA%[1]c[J", 27)

// Writer is a buffered the writer that updates the terminal. The
// contents of writer will be flushed when Flush is called.
type Writer struct {
	out        io.Writer
	buf        bytes.Buffer
	lineCount  int
	fd         uintptr
	isTerminal bool
}

// New returns a new Writer with defaults
func New(out io.Writer) *Writer {
	w := &Writer{out: out}
	if f, ok := out.(*os.File); ok {
		w.fd = f.Fd()
		w.isTerminal = terminal.IsTerminal(int(w.fd))
	}
	return w
}

// Flush flushes the underlying buffer
func (w *Writer) Flush(lineCount int) (err error) {
	if w.lineCount > 0 {
		w.clearLines()
	}
	w.lineCount = lineCount
	_, err = w.buf.WriteTo(w.out)
	return
}

// Write appends the contents of p to the underlying buffer
func (w *Writer) Write(p []byte) (n int, err error) {
	return w.buf.Write(p)
}

// WriteString writes string to the underlying buffer
func (w *Writer) WriteString(s string) (n int, err error) {
	return w.buf.WriteString(s)
}

// ReadFrom reads from the provided io.Reader and writes to the
// underlying buffer.
func (w *Writer) ReadFrom(r io.Reader) (n int64, err error) {
	return w.buf.ReadFrom(r)
}

// GetWidth returns width of underlying terminal.
func (w *Writer) GetWidth() (int, error) {
	if w.isTerminal {
		tw, _, err := terminal.GetSize(int(w.fd))
		return tw, err
	}
	return -1, NotATTY
}
<commit_msg>NotATTY doc comment<commit_after>package cwriter

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"

	"golang.org/x/crypto/ssh/terminal"
)

// NotATTY not a TeleTYpewriter error
var NotATTY = errors.New("not a terminal")

var cuuAndEd = fmt.Sprintf("%c[%%dA%[1]c[J", 27)

// Writer is a buffered the writer that updates the terminal. The
// contents of writer will be flushed when Flush is called.
type Writer struct {
	out        io.Writer
	buf        bytes.Buffer
	lineCount  int
	fd         uintptr
	isTerminal bool
}

// New returns a new Writer with defaults
func New(out io.Writer) *Writer {
	w := &Writer{out: out}
	if f, ok := out.(*os.File); ok {
		w.fd = f.Fd()
		w.isTerminal = terminal.IsTerminal(int(w.fd))
	}
	return w
}

// Flush flushes the underlying buffer
func (w *Writer) Flush(lineCount int) (err error) {
	if w.lineCount > 0 {
		w.clearLines()
	}
	w.lineCount = lineCount
	_, err = w.buf.WriteTo(w.out)
	return
}

// Write appends the contents of p to the underlying buffer
func (w *Writer) Write(p []byte) (n int, err error) {
	return w.buf.Write(p)
}

// WriteString writes string to the underlying buffer
func (w *Writer) WriteString(s string) (n int, err error) {
	return w.buf.WriteString(s)
}

// ReadFrom reads from the provided io.Reader and writes to the
// underlying buffer.
func (w *Writer) ReadFrom(r io.Reader) (n int64, err error) {
	return w.buf.ReadFrom(r)
}

// GetWidth returns width of underlying terminal.
func (w *Writer) GetWidth() (int, error) {
	if w.isTerminal {
		tw, _, err := terminal.GetSize(int(w.fd))
		return tw, err
	}
	return -1, NotATTY
}
<|endoftext|>