<commit_before>package sparkpost

import (
	"encoding/json"
	"errors"
	"net/http"

	"github.com/jmcvetta/napping"
	log "gopkg.in/inconshreveable/log15.v2"

	"github.com/schmooser/go-mailer/message"
)

// Log represents logger.
var Log = log.New()

// SparkPost defines SparkPost transactional mail provider.
type SparkPost struct {
	key string
}

const apiURL = "https://api.sparkpost.com/api/v1/transmissions?num_rcpt_errors=30"

// Address is SparkPost' address part.
type Address struct {
	Name     string `json:"name,omitempty"`
	Email    string `json:"email"`
	HeaderTo string `json:"header_to,omitempty"`
}

// Recipient defines SparkPost' recipient.
type Recipient struct {
	Address Address `json:"address"`
}

// Content defines SparkPost' content.
type Content struct {
	From    Address           `json:"from,omitempty"`
	Headers map[string]string `json:"headers,omitempty"`
	Subject string            `json:"subject,omitempty"`
	Text    string            `json:"text,omitempty"`
	HTML    string            `json:"html,omitempty"`
}

// Message defines SparkPost' message.
type Message struct {
	ReturnPath string      `json:"return_path,omitempty"`
	Recipients []Recipient `json:"recipients,omitempty"`
	Content    Content     `json:"content,omitempt"`
}

// New returns new SparkPost instance. API key is validated upon creation,
// returning error if key is not valid.
func New(key string) (*SparkPost, error) {
	return &SparkPost{key: key}, nil
}

// Send sends provided message in async or sync way.
func (sp *SparkPost) Send(msg *message.Message, async bool) (interface{}, error) {

	var recipients []Recipient
	for _, r := range msg.To {
		recipients = append(recipients, Recipient{
			Address: Address{
				Email: r.Address,
				Name:  r.Name,
			},
		})
	}
	for _, r := range append(msg.CC, msg.BCC...) {
		recipients = append(recipients, Recipient{
			Address: Address{
				Email:    r.Address,
				HeaderTo: msg.To.String(),
			},
		})
	}

	m := Message{
		ReturnPath: msg.From.Address,
		Recipients: recipients,
		Content: Content{
			From: Address{
				Name:  msg.From.Name,
				Email: msg.From.Address,
			},
			Subject: msg.Subject,
			Text:    msg.Text,
			HTML:    msg.HTML,
			Headers: map[string]string{"CC": msg.CC.String()},
		},
	}

	b, _ := json.Marshal(m)
	Log.Debug("Message to send", "msg", string(b))

	headers := make(http.Header)
	headers.Add("Authorization", sp.key)
	s := napping.Session{
		Header: &headers,
	}

	var result interface{}
	resp, err := s.Post(apiURL, &m, &result, nil)
	if err != nil {
		return nil, err
	}

	if resp.Status() == 200 {
		return result, nil
	}

	Log.Error("Response", "status", resp.Status(), "body", resp.RawText())
	return nil, errors.New("Non-200 status returned")
}

func init() {
	Log.SetHandler(log.DiscardHandler())
}
<commit_msg>Fix sparkpost provider - don't generate empty CC header<commit_after>package sparkpost

import (
	"encoding/json"
	"errors"
	"net/http"

	"github.com/jmcvetta/napping"
	log "gopkg.in/inconshreveable/log15.v2"

	"github.com/schmooser/go-mailer/message"
)

// Log represents logger.
var Log = log.New()

// SparkPost defines SparkPost transactional mail provider.
type SparkPost struct {
	key string
}

const apiURL = "https://api.sparkpost.com/api/v1/transmissions?num_rcpt_errors=30"

// Address is SparkPost' address part.
type Address struct {
	Name     string `json:"name,omitempty"`
	Email    string `json:"email"`
	HeaderTo string `json:"header_to,omitempty"`
}

// Recipient defines SparkPost' recipient.
type Recipient struct {
	Address Address `json:"address"`
}

// Content defines SparkPost' content.
type Content struct {
	From    Address           `json:"from,omitempty"`
	Headers map[string]string `json:"headers,omitempty"`
	Subject string            `json:"subject,omitempty"`
	Text    string            `json:"text,omitempty"`
	HTML    string            `json:"html,omitempty"`
}

// Message defines SparkPost' message.
type Message struct {
	ReturnPath string      `json:"return_path,omitempty"`
	Recipients []Recipient `json:"recipients,omitempty"`
	Content    Content     `json:"content,omitempt"`
}

// New returns new SparkPost instance. API key is validated upon creation,
// returning error if key is not valid.
func New(key string) (*SparkPost, error) {
	return &SparkPost{key: key}, nil
}

// Send sends provided message in async or sync way.
func (sp *SparkPost) Send(msg *message.Message, async bool) (interface{}, error) {

	var recipients []Recipient
	for _, r := range msg.To {
		recipients = append(recipients, Recipient{
			Address: Address{
				Email: r.Address,
				Name:  r.Name,
			},
		})
	}
	for _, r := range append(msg.CC, msg.BCC...) {
		recipients = append(recipients, Recipient{
			Address: Address{
				Email:    r.Address,
				HeaderTo: msg.To.String(),
			},
		})
	}

	m := Message{
		ReturnPath: msg.From.Address,
		Recipients: recipients,
		Content: Content{
			From: Address{
				Name:  msg.From.Name,
				Email: msg.From.Address,
			},
			Subject: msg.Subject,
			Text:    msg.Text,
			HTML:    msg.HTML,
			Headers: make(map[string]string),
		},
	}

	if msg.CC.String() != "" {
		m.Content.Headers["CC"] = msg.CC.String()
	}

	b, _ := json.Marshal(m)
	Log.Debug("Message to send", "msg", string(b))

	headers := make(http.Header)
	headers.Add("Authorization", sp.key)
	s := napping.Session{
		Header: &headers,
	}

	var result interface{}
	resp, err := s.Post(apiURL, &m, &result, nil)
	if err != nil {
		return nil, err
	}

	if resp.Status() == 200 {
		return result, nil
	}

	Log.Error("Response", "status", resp.Status(), "body", resp.RawText())
	return nil, errors.New("Non-200 status returned")
}

func init() {
	Log.SetHandler(log.DiscardHandler())
}
<|endoftext|>


<commit_before>/*
Copyright 2019 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 interrupts

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"math/big"
	"net"
	"net/http"
	"os"
	"sync"
	"syscall"
	"testing"
	"time"
)

// interrupt allows for tests to trigger an interrupt as needed
var interrupt = make(chan os.Signal, 1)

// this init will be executed before that in the code package,
// so we can inject our implementation of the interrupt channel
func init() {
	signalsLock.Lock()
	gracePeriod = 0 * time.Second
	signals = func() <-chan os.Signal {
		return interrupt
	}
	signalsLock.Unlock()
}

// instead of building a mechanism to reset/re-initialize the interrupt
// manager which would only be used in testing, we write an integration
// test that only fires the mock interrupt once
func TestInterrupts(t *testing.T) {
	// we need to lock around values used to test otherwise the test
	// goroutine will race with the workers
	lock := sync.Mutex{}

	ctx := Context()
	var ctxDone bool
	go func() {
		<-ctx.Done()

		lock.Lock()
		ctxDone = true
		lock.Unlock()
	}()

	var workDone bool
	var workCancelled bool
	work := func(ctx context.Context) {
		lock.Lock()
		workDone = true
		lock.Unlock()

		<-ctx.Done()

		lock.Lock()
		workCancelled = true
		lock.Unlock()
	}
	Run(work)

	// we cannot use httptest mocks for the tests here as they expect
	// to be started by the httptest package itself, not by a downstream
	// caller like the interrupts library
	var serverCalled bool
	var serverCancelled bool
	server := &http.Server{Addr: ":9999", Handler: http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
		lock.Lock()
		serverCalled = true
		lock.Unlock()
	})}
	server.RegisterOnShutdown(func() {
		lock.Lock()
		serverCancelled = true
		lock.Unlock()
	})
	ListenAndServe(server, 1*time.Microsecond)
	if _, err := http.Get("http://127.0.0.1:9999"); err != nil {
		t.Errorf("could not reach server registered with ListenAndServe(): %v", err)
	}

	var tlsServerCalled bool
	var tlsServerCancelled bool
	tlsServer := &http.Server{Addr: "127.0.0.1:9998", Handler: http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
		lock.Lock()
		tlsServerCalled = true
		lock.Unlock()
	})}
	tlsServer.RegisterOnShutdown(func() {
		lock.Lock()
		tlsServerCancelled = true
		lock.Unlock()
	})
	cert, key, err := generateCerts("127.0.0.1")
	if err != nil {
		t.Fatalf("could not generate cert and key for TLS server: %v", err)
	}
	ListenAndServeTLS(tlsServer, cert, key, 1*time.Microsecond)
	client := &http.Client{Transport: &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: true}}}
	if _, err := client.Get("https://127.0.0.1:9998"); err != nil {
		t.Errorf("could not reach server registered with ListenAndServeTLS(): %v", err)
	}

	var intervalCalls int
	interval := func() time.Duration {
		lock.Lock()
		intervalCalls++
		lock.Unlock()
		if intervalCalls > 2 {
			return 10 * time.Hour
		}
		return 1 * time.Nanosecond
	}
	var tickCalls int
	tick := func() {
		lock.Lock()
		tickCalls++
		lock.Unlock()
	}
	Tick(tick, interval)
	// writing a test that functions correctly here without being susceptible
	// to timing flakes is challenging. Using time.Sleep like this does have
	// that downside, but the sleep time is many orders of magnitude higher
	// than the tick intervals and the amount of time taken to execute the
	// test as well, so it is going to be exceedingly rare that scheduling of
	// the test process will cause a flake here from timing. The test cannot
	// use synchronized approaches to waiting here as we do not know how long
	// we must wait. The test must have enough time to ask for the interval
	// as many times as we expect it to, but if we only wait for that we fail
	// to catch the cases where the interval is requested too many times.
	time.Sleep(100 * time.Millisecond)

	var onInterruptCalled bool
	OnInterrupt(func() {
		lock.Lock()
		onInterruptCalled = true
		lock.Unlock()
	})

	done := sync.WaitGroup{}
	done.Add(1)
	go func() {
		WaitForGracefulShutdown()
		done.Done()
	}()

	if onInterruptCalled {
		t.Error("work registered with OnInterrupt() was executed before interrupt")
	}

	// trigger the interrupt
	interrupt <- syscall.Signal(1)
	// wait for graceful shutdown to occur
	done.Wait()

	lock.Lock()
	if !ctxDone {
		t.Error("context from Context() was not cancelled on interrupt")
	}
	if !workDone {
		t.Error("work registered with Run() was not executed")
	}
	if !workCancelled {
		t.Error("work registered with Run() was not cancelled on interrupt")
	}
	if !serverCalled {
		t.Error("server registered with ListenAndServe() was not serving")
	}
	if !serverCancelled {
		t.Error("server registered with ListenAndServe() was not cancelled on interrupt")
	}
	if !tlsServerCalled {
		t.Error("server registered with ListenAndServeTLS() was not serving")
	}
	if !tlsServerCancelled {
		t.Error("server registered with ListenAndServeTLS() was not cancelled on interrupt")
	}
	if tickCalls != 2 {
		t.Errorf("work registered with Tick() was called %d times, not %d; interval was requested %d times", tickCalls, 2, intervalCalls)
	}
	if !onInterruptCalled {
		t.Error("work registered with OnInterrupt() was not executed on interrupt")
	}
	lock.Unlock()
}

func generateCerts(url string) (string, string, error) {
	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return "", "", fmt.Errorf("failed to generate private key: %v", err)
	}

	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		return "", "", fmt.Errorf("failed to generate serial number: %s", err)
	}

	template := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			Organization: []string{"Acme Co"},
		},
		NotBefore: time.Now(),
		NotAfter:  time.Now().Add(1 * time.Hour),

		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,

		IPAddresses: []net.IP{net.ParseIP(url)},
	}

	derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, &priv.PublicKey, priv)
	if err != nil {
		return "", "", fmt.Errorf("failed to create certificate: %s", err)
	}

	certOut, err := ioutil.TempFile("", "cert.pem")
	if err != nil {
		return "", "", fmt.Errorf("failed to open cert.pem for writing: %s", err)
	}
	if err := pem.Encode(certOut, &pem.Block{Type: "CERTIFICATE", Bytes: derBytes}); err != nil {
		return "", "", fmt.Errorf("failed to write data to cert.pem: %s", err)
	}
	if err := certOut.Close(); err != nil {
		return "", "", fmt.Errorf("error closing cert.pem: %s", err)
	}

	keyOut, err := ioutil.TempFile("", "key.pem")
	if err != nil {
		return "", "", fmt.Errorf("failed to open key.pem for writing: %v", err)
	}
	privBytes, err := x509.MarshalPKCS8PrivateKey(priv)
	if err != nil {
		return "", "", fmt.Errorf("unable to marshal private key: %v", err)
	}
	if err := pem.Encode(keyOut, &pem.Block{Type: "PRIVATE KEY", Bytes: privBytes}); err != nil {
		return "", "", fmt.Errorf("failed to write data to key.pem: %s", err)
	}
	if err := keyOut.Close(); err != nil {
		return "", "", fmt.Errorf("error closing key.pem: %s", err)
	}
	if err := os.Chmod(keyOut.Name(), 0600); err != nil {
		return "", "", fmt.Errorf("could not change permissions on key.pem: %v", err)
	}
	return certOut.Name(), keyOut.Name(), nil
}
<commit_msg>Update interrupts test servers to grab random ports<commit_after>/*
Copyright 2019 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 interrupts

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"math/big"
	"net"
	"net/http"
	"os"
	"sync"
	"syscall"
	"testing"
	"time"
)

// interrupt allows for tests to trigger an interrupt as needed
var interrupt = make(chan os.Signal, 1)

// this init will be executed before that in the code package,
// so we can inject our implementation of the interrupt channel
func init() {
	signalsLock.Lock()
	gracePeriod = 0 * time.Second
	signals = func() <-chan os.Signal {
		return interrupt
	}
	signalsLock.Unlock()
}

// instead of building a mechanism to reset/re-initialize the interrupt
// manager which would only be used in testing, we write an integration
// test that only fires the mock interrupt once
func TestInterrupts(t *testing.T) {
	// we need to lock around values used to test otherwise the test
	// goroutine will race with the workers
	lock := sync.Mutex{}

	ctx := Context()
	var ctxDone bool
	go func() {
		<-ctx.Done()

		lock.Lock()
		ctxDone = true
		lock.Unlock()
	}()

	var workDone bool
	var workCancelled bool
	work := func(ctx context.Context) {
		lock.Lock()
		workDone = true
		lock.Unlock()

		<-ctx.Done()

		lock.Lock()
		workCancelled = true
		lock.Unlock()
	}
	Run(work)

	// we cannot use httptest mocks for the tests here as they expect
	// to be started by the httptest package itself, not by a downstream
	// caller like the interrupts library
	var serverCalled bool
	var serverCancelled bool
	listener, err := net.Listen("tcp", "127.0.0.1:")
	if err != nil {
		t.Fatalf("could not listen on random port: %v", err)
	}
	if err := listener.Close(); err != nil {
		t.Fatalf("could close listener: %v", err)
	}
	server := &http.Server{Addr: listener.Addr().String(), Handler: http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
		lock.Lock()
		serverCalled = true
		lock.Unlock()
	})}
	server.RegisterOnShutdown(func() {
		lock.Lock()
		serverCancelled = true
		lock.Unlock()
	})
	ListenAndServe(server, 1*time.Microsecond)
	if _, err := http.Get("http://" + listener.Addr().String()); err != nil {
		t.Errorf("could not reach server registered with ListenAndServe(): %v", err)
	}

	var tlsServerCalled bool
	var tlsServerCancelled bool
	tlsListener, err := net.Listen("tcp", "127.0.0.1:")
	if err != nil {
		t.Fatalf("could not listen on random port: %v", err)
	}
	if err := tlsListener.Close(); err != nil {
		t.Fatalf("could close listener: %v", err)
	}
	tlsServer := &http.Server{Addr: tlsListener.Addr().String(), Handler: http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
		lock.Lock()
		tlsServerCalled = true
		lock.Unlock()
	})}
	tlsServer.RegisterOnShutdown(func() {
		lock.Lock()
		tlsServerCancelled = true
		lock.Unlock()
	})
	cert, key, err := generateCerts("127.0.0.1")
	if err != nil {
		t.Fatalf("could not generate cert and key for TLS server: %v", err)
	}
	ListenAndServeTLS(tlsServer, cert, key, 1*time.Microsecond)
	client := &http.Client{Transport: &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: true}}}
	if _, err := client.Get("https://" + tlsListener.Addr().String()); err != nil {
		t.Errorf("could not reach server registered with ListenAndServeTLS(): %v", err)
	}

	var intervalCalls int
	interval := func() time.Duration {
		lock.Lock()
		intervalCalls++
		lock.Unlock()
		if intervalCalls > 2 {
			return 10 * time.Hour
		}
		return 1 * time.Nanosecond
	}
	var tickCalls int
	tick := func() {
		lock.Lock()
		tickCalls++
		lock.Unlock()
	}
	Tick(tick, interval)
	// writing a test that functions correctly here without being susceptible
	// to timing flakes is challenging. Using time.Sleep like this does have
	// that downside, but the sleep time is many orders of magnitude higher
	// than the tick intervals and the amount of time taken to execute the
	// test as well, so it is going to be exceedingly rare that scheduling of
	// the test process will cause a flake here from timing. The test cannot
	// use synchronized approaches to waiting here as we do not know how long
	// we must wait. The test must have enough time to ask for the interval
	// as many times as we expect it to, but if we only wait for that we fail
	// to catch the cases where the interval is requested too many times.
	time.Sleep(100 * time.Millisecond)

	var onInterruptCalled bool
	OnInterrupt(func() {
		lock.Lock()
		onInterruptCalled = true
		lock.Unlock()
	})

	done := sync.WaitGroup{}
	done.Add(1)
	go func() {
		WaitForGracefulShutdown()
		done.Done()
	}()

	if onInterruptCalled {
		t.Error("work registered with OnInterrupt() was executed before interrupt")
	}

	// trigger the interrupt
	interrupt <- syscall.Signal(1)
	// wait for graceful shutdown to occur
	done.Wait()

	lock.Lock()
	if !ctxDone {
		t.Error("context from Context() was not cancelled on interrupt")
	}
	if !workDone {
		t.Error("work registered with Run() was not executed")
	}
	if !workCancelled {
		t.Error("work registered with Run() was not cancelled on interrupt")
	}
	if !serverCalled {
		t.Error("server registered with ListenAndServe() was not serving")
	}
	if !serverCancelled {
		t.Error("server registered with ListenAndServe() was not cancelled on interrupt")
	}
	if !tlsServerCalled {
		t.Error("server registered with ListenAndServeTLS() was not serving")
	}
	if !tlsServerCancelled {
		t.Error("server registered with ListenAndServeTLS() was not cancelled on interrupt")
	}
	if tickCalls != 2 {
		t.Errorf("work registered with Tick() was called %d times, not %d; interval was requested %d times", tickCalls, 2, intervalCalls)
	}
	if !onInterruptCalled {
		t.Error("work registered with OnInterrupt() was not executed on interrupt")
	}
	lock.Unlock()
}

func generateCerts(url string) (string, string, error) {
	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return "", "", fmt.Errorf("failed to generate private key: %v", err)
	}

	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		return "", "", fmt.Errorf("failed to generate serial number: %s", err)
	}

	template := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			Organization: []string{"Acme Co"},
		},
		NotBefore: time.Now(),
		NotAfter:  time.Now().Add(1 * time.Hour),

		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,

		IPAddresses: []net.IP{net.ParseIP(url)},
	}

	derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, &priv.PublicKey, priv)
	if err != nil {
		return "", "", fmt.Errorf("failed to create certificate: %s", err)
	}

	certOut, err := ioutil.TempFile("", "cert.pem")
	if err != nil {
		return "", "", fmt.Errorf("failed to open cert.pem for writing: %s", err)
	}
	if err := pem.Encode(certOut, &pem.Block{Type: "CERTIFICATE", Bytes: derBytes}); err != nil {
		return "", "", fmt.Errorf("failed to write data to cert.pem: %s", err)
	}
	if err := certOut.Close(); err != nil {
		return "", "", fmt.Errorf("error closing cert.pem: %s", err)
	}

	keyOut, err := ioutil.TempFile("", "key.pem")
	if err != nil {
		return "", "", fmt.Errorf("failed to open key.pem for writing: %v", err)
	}
	privBytes, err := x509.MarshalPKCS8PrivateKey(priv)
	if err != nil {
		return "", "", fmt.Errorf("unable to marshal private key: %v", err)
	}
	if err := pem.Encode(keyOut, &pem.Block{Type: "PRIVATE KEY", Bytes: privBytes}); err != nil {
		return "", "", fmt.Errorf("failed to write data to key.pem: %s", err)
	}
	if err := keyOut.Close(); err != nil {
		return "", "", fmt.Errorf("error closing key.pem: %s", err)
	}
	if err := os.Chmod(keyOut.Name(), 0600); err != nil {
		return "", "", fmt.Errorf("could not change permissions on key.pem: %v", err)
	}
	return certOut.Name(), keyOut.Name(), nil
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"sync"

	"github.com/Microsoft/hcsshim/internal/oci"
	"github.com/Microsoft/hcsshim/internal/uvm"
	eventstypes "github.com/containerd/containerd/api/events"
	"github.com/containerd/containerd/errdefs"
	"github.com/containerd/containerd/runtime"
	"github.com/containerd/containerd/runtime/v2/task"
	specs "github.com/opencontainers/runtime-spec/specs-go"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"golang.org/x/sync/errgroup"
)

// shimPod represents the logical grouping of all tasks in a single set of
// shared namespaces. The pod sandbox (container) is represented by the task
// that matches the `shimPod.ID()`
type shimPod interface {
	// ID is the id of the task representing the pause (sandbox) container.
	ID() string
	// CreateTask creates a workload task within this pod named `tid` with
	// settings `s`.
	//
	// If `tid==ID()` or `tid` is the same as any other task in this pod, this
	// pod MUST return `errdefs.ErrAlreadyExists`.
	CreateTask(ctx context.Context, req *task.CreateTaskRequest, s *specs.Spec) (shimTask, error)
	// GetTask returns a task in this pod that matches `tid`.
	//
	// If `tid` is not found, this pod MUST return `errdefs.ErrNotFound`.
	GetTask(tid string) (shimTask, error)
	// KillTask sends `signal` to task that matches `tid`.
	//
	// If `tid` is not found, this pod MUST return `errdefs.ErrNotFound`.
	//
	// If `tid==ID() && eid == "" && all == true` this pod will send `signal` to
	// all tasks in the pod and lastly send `signal` to the sandbox itself.
	//
	// If `all == true && eid != ""` this pod MUST return
	// `errdefs.ErrFailedPrecondition`.
	//
	// A call to `KillTask` is only valid when the exec found by `tid,eid` is in
	// the `shimExecStateRunning` state. If the exec is not in this state this
	// pod MUST return `errdefs.ErrFailedPrecondition`.
	KillTask(ctx context.Context, tid, eid string, signal uint32, all bool) error
}

func createPod(ctx context.Context, events publisher, req *task.CreateTaskRequest, s *specs.Spec) (shimPod, error) {
	logrus.WithFields(logrus.Fields{
		"tid": req.ID,
	}).Debug("createPod")

	ct, sid, err := oci.GetSandboxTypeAndID(s.Annotations)
	if err != nil {
		return nil, err
	}
	if ct != oci.KubernetesContainerTypeSandbox {
		return nil, errors.Wrapf(
			errdefs.ErrFailedPrecondition,
			"expected annotation: '%s': '%s' got '%s'",
			oci.KubernetesContainerTypeAnnotation,
			oci.KubernetesContainerTypeSandbox,
			ct)
	}
	if sid != req.ID {
		return nil, errors.Wrapf(
			errdefs.ErrFailedPrecondition,
			"expected annotation '%s': '%s' got '%s'",
			oci.KubernetesSandboxIDAnnotation,
			req.ID,
			sid)
	}

	owner, err := os.Executable()
	if err != nil {
		return nil, err
	}

	var parent *uvm.UtilityVM
	if oci.IsIsolated(s) {
		// Create the UVM parent
		opts, err := oci.SpecToUVMCreateOpts(s, fmt.Sprintf("%s@vm", req.ID), owner)
		if err != nil {
			return nil, err
		}
		switch opts.(type) {
		case *uvm.OptionsLCOW:
			lopts := (opts).(*uvm.OptionsLCOW)
			parent, err = uvm.CreateLCOW(lopts)
			if err != nil {
				return nil, err
			}
		case *uvm.OptionsWCOW:
			wopts := (opts).(*uvm.OptionsWCOW)

			// In order for the UVM sandbox.vhdx not to collide with the actual
			// nested Argon sandbox.vhdx we append the \vm folder to the last
			// entry in the list.
			layersLen := len(s.Windows.LayerFolders)
			layers := make([]string, layersLen)
			copy(layers, s.Windows.LayerFolders)

			vmPath := filepath.Join(layers[layersLen-1], "vm")
			err := os.MkdirAll(vmPath, 0)
			if err != nil {
				return nil, err
			}
			layers[layersLen-1] = vmPath
			wopts.LayerFolders = layers

			parent, err = uvm.CreateWCOW(wopts)
			if err != nil {
				return nil, err
			}
		}
		err = parent.Start()
		if err != nil {
			parent.Close()
		}
	} else if !oci.IsWCOW(s) {
		return nil, errors.Wrap(errdefs.ErrFailedPrecondition, "oci spec does not contain WCOW or LCOW spec")
	}

	p := pod{
		events: events,
		id:     req.ID,
		host:   parent,
	}
	if oci.IsWCOW(s) {
		// For WCOW we fake out the init task since we dont need it.
		p.sandboxTask = newWcowPodSandboxTask(ctx, events, req.ID, req.Bundle, parent)
		// Publish the created event. We only do this for the fake WCOW task a
		// HCS Task will event itself.
		events(
			runtime.TaskCreateEventTopic,
			&eventstypes.TaskCreate{
				ContainerID: req.ID,
				Bundle:      req.Bundle,
				Rootfs:      req.Rootfs,
				IO: &eventstypes.TaskIO{
					Stdin:    req.Stdin,
					Stdout:   req.Stdout,
					Stderr:   req.Stderr,
					Terminal: req.Terminal,
				},
				Checkpoint: "",
				Pid:        0,
			})
	} else {
		// LCOW requires a real task for the sandbox
		lt, err := newHcsTask(ctx, events, parent, true, req, s)
		if err != nil {
			parent.Close()
			return nil, err
		}
		p.sandboxTask = lt
	}

	return &p, nil
}

var _ = (shimPod)(&pod{})

type pod struct {
	events publisher
	// id is the id of the sandbox task when the pod is created.
	//
	// It MUST be treated as read only in the lifetime of the pod.
	id string
	// sandboxTask is the task that represents the sandbox.
	//
	// Note: The invariant `id==sandboxTask.ID()` MUST be true.
	//
	// It MUST be treated as read only in the lifetime of the pod.
	sandboxTask shimTask
	// host is the UtilityVM that is hosting `sandboxTask` if the task is
	// hypervisor isolated.
	//
	// It MUST be treated as read only in the lifetime of the pod.
	host *uvm.UtilityVM

	// wcl is the worload create mutex. All calls to CreateTask must hold this
	// lock while the ID reservation takes place. Once the ID is held it is safe
	// to release the lock to allow concurrent creates.
	wcl           sync.Mutex
	workloadTasks sync.Map
}

func (p *pod) ID() string {
	return p.id
}

func (p *pod) CreateTask(ctx context.Context, req *task.CreateTaskRequest, s *specs.Spec) (shimTask, error) {
	logrus.WithFields(logrus.Fields{
		"pod-id": p.id,
		"tid":    req.ID,
	}).Debug("pod::CreateTask")

	if req.ID == p.id {
		return nil, errors.Wrapf(errdefs.ErrAlreadyExists, "task with id: '%s' already exists", req.ID)
	}
	e, _ := p.sandboxTask.GetExec("")
	if e.State() != shimExecStateRunning {
		return nil, errors.Wrapf(errdefs.ErrFailedPrecondition, "task with id: '%s' cannot be created in pod: '%s' which is not running", req.ID, p.id)
	}

	p.wcl.Lock()
	_, loaded := p.workloadTasks.LoadOrStore(req.ID, nil)
	if loaded {
		return nil, errors.Wrapf(errdefs.ErrAlreadyExists, "task with id: '%s' already exists id pod: '%s'", req.ID, p.id)
	}
	p.wcl.Unlock()
	var err error
	defer func() {
		if err != nil {
			p.workloadTasks.Delete(req.ID)
		}
	}()

	ct, sid, err := oci.GetSandboxTypeAndID(s.Annotations)
	if err != nil {
		return nil, err
	}
	if ct != oci.KubernetesContainerTypeContainer {
		return nil, errors.Wrapf(
			errdefs.ErrFailedPrecondition,
			"expected annotation: '%s': '%s' got '%s'",
			oci.KubernetesContainerTypeAnnotation,
			oci.KubernetesContainerTypeContainer,
			ct)
	}
	if sid != p.id {
		return nil, errors.Wrapf(
			errdefs.ErrFailedPrecondition,
			"expected annotation '%s': '%s' got '%s'",
			oci.KubernetesSandboxIDAnnotation,
			p.id,
			sid)
	}

	st, err := newHcsTask(ctx, p.events, p.host, false, req, s)
	if err != nil {
		return nil, err
	}

	p.workloadTasks.Store(req.ID, st)
	return st, nil
}

func (p *pod) GetTask(tid string) (shimTask, error) {
	if tid == p.id {
		return p.sandboxTask, nil
	}
	raw, loaded := p.workloadTasks.Load(tid)
	if !loaded {
		return nil, errors.Wrapf(errdefs.ErrNotFound, "task with id: '%s' not found", tid)
	}
	return raw.(shimTask), nil
}

func (p *pod) KillTask(ctx context.Context, tid, eid string, signal uint32, all bool) error {
	logrus.WithFields(logrus.Fields{
		"pod-id": p.id,
		"tid":    tid,
		"eid":    eid,
		"signal": signal,
		"all":    all,
	}).Debug("pod::KillTask")

	t, err := p.GetTask(tid)
	if err != nil {
		return err
	}
	if all && eid != "" {
		return errors.Wrapf(errdefs.ErrFailedPrecondition, "cannot signal all with non empty ExecID: '%s'", eid)
	}
	eg := errgroup.Group{}
	if all && tid == p.id {
		// We are in a kill all on the sandbox task. Signal everything.
		p.workloadTasks.Range(func(key, value interface{}) bool {
			wt := value.(shimTask)
			eg.Go(func() error {
				return wt.KillExec(ctx, eid, signal, all)
			})

			// iterate all
			return false
		})
	}
	eg.Go(func() error {
		return t.KillExec(ctx, eid, signal, all)
	})
	return eg.Wait()
}
<commit_msg>Pod Kill all should skip exited tasks<commit_after>package main

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"sync"

	"github.com/Microsoft/hcsshim/internal/oci"
	"github.com/Microsoft/hcsshim/internal/uvm"
	eventstypes "github.com/containerd/containerd/api/events"
	"github.com/containerd/containerd/errdefs"
	"github.com/containerd/containerd/runtime"
	"github.com/containerd/containerd/runtime/v2/task"
	specs "github.com/opencontainers/runtime-spec/specs-go"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"golang.org/x/sync/errgroup"
)

// shimPod represents the logical grouping of all tasks in a single set of
// shared namespaces. The pod sandbox (container) is represented by the task
// that matches the `shimPod.ID()`
type shimPod interface {
	// ID is the id of the task representing the pause (sandbox) container.
	ID() string
	// CreateTask creates a workload task within this pod named `tid` with
	// settings `s`.
	//
	// If `tid==ID()` or `tid` is the same as any other task in this pod, this
	// pod MUST return `errdefs.ErrAlreadyExists`.
	CreateTask(ctx context.Context, req *task.CreateTaskRequest, s *specs.Spec) (shimTask, error)
	// GetTask returns a task in this pod that matches `tid`.
	//
	// If `tid` is not found, this pod MUST return `errdefs.ErrNotFound`.
	GetTask(tid string) (shimTask, error)
	// KillTask sends `signal` to task that matches `tid`.
	//
	// If `tid` is not found, this pod MUST return `errdefs.ErrNotFound`.
	//
	// If `tid==ID() && eid == "" && all == true` this pod will send `signal` to
	// all tasks in the pod and lastly send `signal` to the sandbox itself.
	//
	// If `all == true && eid != ""` this pod MUST return
	// `errdefs.ErrFailedPrecondition`.
	//
	// A call to `KillTask` is only valid when the exec found by `tid,eid` is in
	// the `shimExecStateRunning` state. If the exec is not in this state this
	// pod MUST return `errdefs.ErrFailedPrecondition`.
	KillTask(ctx context.Context, tid, eid string, signal uint32, all bool) error
}

func createPod(ctx context.Context, events publisher, req *task.CreateTaskRequest, s *specs.Spec) (shimPod, error) {
	logrus.WithFields(logrus.Fields{
		"tid": req.ID,
	}).Debug("createPod")

	ct, sid, err := oci.GetSandboxTypeAndID(s.Annotations)
	if err != nil {
		return nil, err
	}
	if ct != oci.KubernetesContainerTypeSandbox {
		return nil, errors.Wrapf(
			errdefs.ErrFailedPrecondition,
			"expected annotation: '%s': '%s' got '%s'",
			oci.KubernetesContainerTypeAnnotation,
			oci.KubernetesContainerTypeSandbox,
			ct)
	}
	if sid != req.ID {
		return nil, errors.Wrapf(
			errdefs.ErrFailedPrecondition,
			"expected annotation '%s': '%s' got '%s'",
			oci.KubernetesSandboxIDAnnotation,
			req.ID,
			sid)
	}

	owner, err := os.Executable()
	if err != nil {
		return nil, err
	}

	var parent *uvm.UtilityVM
	if oci.IsIsolated(s) {
		// Create the UVM parent
		opts, err := oci.SpecToUVMCreateOpts(s, fmt.Sprintf("%s@vm", req.ID), owner)
		if err != nil {
			return nil, err
		}
		switch opts.(type) {
		case *uvm.OptionsLCOW:
			lopts := (opts).(*uvm.OptionsLCOW)
			parent, err = uvm.CreateLCOW(lopts)
			if err != nil {
				return nil, err
			}
		case *uvm.OptionsWCOW:
			wopts := (opts).(*uvm.OptionsWCOW)

			// In order for the UVM sandbox.vhdx not to collide with the actual
			// nested Argon sandbox.vhdx we append the \vm folder to the last
			// entry in the list.
			layersLen := len(s.Windows.LayerFolders)
			layers := make([]string, layersLen)
			copy(layers, s.Windows.LayerFolders)

			vmPath := filepath.Join(layers[layersLen-1], "vm")
			err := os.MkdirAll(vmPath, 0)
			if err != nil {
				return nil, err
			}
			layers[layersLen-1] = vmPath
			wopts.LayerFolders = layers

			parent, err = uvm.CreateWCOW(wopts)
			if err != nil {
				return nil, err
			}
		}
		err = parent.Start()
		if err != nil {
			parent.Close()
		}
	} else if !oci.IsWCOW(s) {
		return nil, errors.Wrap(errdefs.ErrFailedPrecondition, "oci spec does not contain WCOW or LCOW spec")
	}

	p := pod{
		events: events,
		id:     req.ID,
		host:   parent,
	}
	if oci.IsWCOW(s) {
		// For WCOW we fake out the init task since we dont need it.
		p.sandboxTask = newWcowPodSandboxTask(ctx, events, req.ID, req.Bundle, parent)
		// Publish the created event. We only do this for the fake WCOW task a
		// HCS Task will event itself.
		events(
			runtime.TaskCreateEventTopic,
			&eventstypes.TaskCreate{
				ContainerID: req.ID,
				Bundle:      req.Bundle,
				Rootfs:      req.Rootfs,
				IO: &eventstypes.TaskIO{
					Stdin:    req.Stdin,
					Stdout:   req.Stdout,
					Stderr:   req.Stderr,
					Terminal: req.Terminal,
				},
				Checkpoint: "",
				Pid:        0,
			})
	} else {
		// LCOW requires a real task for the sandbox
		lt, err := newHcsTask(ctx, events, parent, true, req, s)
		if err != nil {
			parent.Close()
			return nil, err
		}
		p.sandboxTask = lt
	}

	return &p, nil
}

var _ = (shimPod)(&pod{})

type pod struct {
	events publisher
	// id is the id of the sandbox task when the pod is created.
	//
	// It MUST be treated as read only in the lifetime of the pod.
	id string
	// sandboxTask is the task that represents the sandbox.
	//
	// Note: The invariant `id==sandboxTask.ID()` MUST be true.
	//
	// It MUST be treated as read only in the lifetime of the pod.
	sandboxTask shimTask
	// host is the UtilityVM that is hosting `sandboxTask` if the task is
	// hypervisor isolated.
	//
	// It MUST be treated as read only in the lifetime of the pod.
	host *uvm.UtilityVM

	// wcl is the worload create mutex. All calls to CreateTask must hold this
	// lock while the ID reservation takes place. Once the ID is held it is safe
	// to release the lock to allow concurrent creates.
	wcl           sync.Mutex
	workloadTasks sync.Map
}

func (p *pod) ID() string {
	return p.id
}

func (p *pod) CreateTask(ctx context.Context, req *task.CreateTaskRequest, s *specs.Spec) (shimTask, error) {
	logrus.WithFields(logrus.Fields{
		"pod-id": p.id,
		"tid":    req.ID,
	}).Debug("pod::CreateTask")

	if req.ID == p.id {
		return nil, errors.Wrapf(errdefs.ErrAlreadyExists, "task with id: '%s' already exists", req.ID)
	}
	e, _ := p.sandboxTask.GetExec("")
	if e.State() != shimExecStateRunning {
		return nil, errors.Wrapf(errdefs.ErrFailedPrecondition, "task with id: '%s' cannot be created in pod: '%s' which is not running", req.ID, p.id)
	}

	p.wcl.Lock()
	_, loaded := p.workloadTasks.LoadOrStore(req.ID, nil)
	if loaded {
		return nil, errors.Wrapf(errdefs.ErrAlreadyExists, "task with id: '%s' already exists id pod: '%s'", req.ID, p.id)
	}
	p.wcl.Unlock()
	var err error
	defer func() {
		if err != nil {
			p.workloadTasks.Delete(req.ID)
		}
	}()

	ct, sid, err := oci.GetSandboxTypeAndID(s.Annotations)
	if err != nil {
		return nil, err
	}
	if ct != oci.KubernetesContainerTypeContainer {
		return nil, errors.Wrapf(
			errdefs.ErrFailedPrecondition,
			"expected annotation: '%s': '%s' got '%s'",
			oci.KubernetesContainerTypeAnnotation,
			oci.KubernetesContainerTypeContainer,
			ct)
	}
	if sid != p.id {
		return nil, errors.Wrapf(
			errdefs.ErrFailedPrecondition,
			"expected annotation '%s': '%s' got '%s'",
			oci.KubernetesSandboxIDAnnotation,
			p.id,
			sid)
	}

	st, err := newHcsTask(ctx, p.events, p.host, false, req, s)
	if err != nil {
		return nil, err
	}

	p.workloadTasks.Store(req.ID, st)
	return st, nil
}

func (p *pod) GetTask(tid string) (shimTask, error) {
	if tid == p.id {
		return p.sandboxTask, nil
	}
	raw, loaded := p.workloadTasks.Load(tid)
	if !loaded {
		return nil, errors.Wrapf(errdefs.ErrNotFound, "task with id: '%s' not found", tid)
	}
	return raw.(shimTask), nil
}

func (p *pod) KillTask(ctx context.Context, tid, eid string, signal uint32, all bool) error {
	logrus.WithFields(logrus.Fields{
		"pod-id": p.id,
		"tid":    tid,
		"eid":    eid,
		"signal": signal,
		"all":    all,
	}).Debug("pod::KillTask")

	t, err := p.GetTask(tid)
	if err != nil {
		return err
	}
	if all && eid != "" {
		return errors.Wrapf(errdefs.ErrFailedPrecondition, "cannot signal all with non empty ExecID: '%s'", eid)
	}
	eg := errgroup.Group{}
	if all && tid == p.id {
		// We are in a kill all on the sandbox task. Signal everything.
		p.workloadTasks.Range(func(key, value interface{}) bool {
			wt := value.(shimTask)
			ie, _ := wt.GetExec("")
			// Only send the kill signal to non-exited tasks when in the `all`
			// case.
			if ie.State() != shimExecStateExited {
				eg.Go(func() error {
					return wt.KillExec(ctx, eid, signal, all)
				})
			}

			// iterate all
			return false
		})
	}
	eg.Go(func() error {
		return t.KillExec(ctx, eid, signal, all)
	})
	return eg.Wait()
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strconv"
	"strings"
	"sync"
	"time"

	//"runtime/pprof"

	"github.com/golang/glog"
	"github.com/livepeer/go-livepeer/common"
	"github.com/livepeer/lpms/ffmpeg"
	"github.com/livepeer/m3u8"
	"github.com/olekukonko/tablewriter"
)

func main() {
	// Override the default flag set since there are dependencies that
	// incorrectly add their own flags (specifically, due to the 'testing'
	// package being linked)
	flag.Set("logtostderr", "true")
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)

	in := flag.String("in", "", "Input m3u8 manifest file")
	live := flag.Bool("live", true, "Simulate live stream")
	concurrentSessions := flag.Int("concurrentSessions", 1, "# of concurrent transcode sessions")
	segs := flag.Int("segs", 0, "Maximum # of segments to transcode (default all)")
	transcodingOptions := flag.String("transcodingOptions", "P240p30fps16x9,P360p30fps16x9,P720p30fps16x9", "Transcoding options for broadcast job, or path to json config")
	nvidia := flag.String("nvidia", "", "Comma-separated list of Nvidia GPU device IDs (or \"all\" for all available devices)")
	outPrefix := flag.String("outPrefix", "", "Output segments' prefix (no segments are generated by default)")

	flag.Parse()

	if *in == "" {
		glog.Errorf("Please provide the input manifest as `%s -in <input.m3u8>`", os.Args[0])
		flag.Usage()
		os.Exit(1)
	}

	profiles := parseVideoProfiles(*transcodingOptions)

	f, err := os.Open(*in)
	if err != nil {
		glog.Fatal("Couldn't open input manifest: ", err)
	}
	p, _, err := m3u8.DecodeFrom(bufio.NewReader(f), true)
	if err != nil {
		glog.Fatal("Couldn't decode input manifest: ", err)
	}
	pl, ok := p.(*m3u8.MediaPlaylist)
	if !ok {
		glog.Fatalf("Expecting media playlist in the input %s", *in)
	}

	accel := ffmpeg.Software
	devices := []string{}
	if *nvidia != "" {
		var err error
		accel = ffmpeg.Nvidia
		devices, err = common.ParseNvidiaDevices(*nvidia)
		if err != nil {
			glog.Fatalf("Error while parsing '-nvidia %v' flag: %v", *nvidia, err)
		}
	}

	ffmpeg.InitFFmpeg()
	var wg sync.WaitGroup
	dir := path.Dir(*in)

	table := tablewriter.NewWriter(os.Stderr)
	data := [][]string{
		{"Source File", *in},
		{"Transcoding Options", *transcodingOptions},
		{"Concurrent Sessions", fmt.Sprintf("%v", *concurrentSessions)},
		{"Live Mode", fmt.Sprintf("%v", *live)},
	}

	if accel == ffmpeg.Nvidia && len(devices) > 0 {
		data = append(data, []string{"Nvidia GPU IDs", fmt.Sprintf("%v", strings.Join(devices, ","))})
	}

	table.SetAlignment(tablewriter.ALIGN_LEFT)
	table.SetCenterSeparator("*")
	table.SetColumnSeparator("|")
	table.AppendBulk(data)
	table.Render()

	fmt.Println("timestamp,session,segment,seg_dur,transcode_time")
	segCount := 0
	realTimeSegCount := 0
	srcDur := 0.0
	var mu sync.Mutex
	transcodeDur := 0.0
	for i := 0; i < *concurrentSessions; i++ {
		wg.Add(1)
		go func(k int, wg *sync.WaitGroup) {
			tc := ffmpeg.NewTranscoder()
			for j, v := range pl.Segments {
				iterStart := time.Now()
				if *segs > 0 && j >= *segs {
					break
				}
				if v == nil {
					continue
				}
				u := path.Join(dir, v.URI)
				in := &ffmpeg.TranscodeOptionsIn{
					Fname: u,
					Accel: accel,
				}
				if ffmpeg.Software != accel {
					in.Device = devices[k%len(devices)]
				}
				profs2opts := func(profs []ffmpeg.VideoProfile) []ffmpeg.TranscodeOptions {
					opts := []ffmpeg.TranscodeOptions{}
					for n, p := range profs {
						oname := ""
						muxer := ""
						if *outPrefix != "" {
							oname = fmt.Sprintf("%s_%s_%d_%d_%d.ts", *outPrefix, p.Name, n, k, j)
							muxer = "mpegts"
						} else {
							oname = "-"
							muxer = "null"
						}
						o := ffmpeg.TranscodeOptions{
							Oname:        oname,
							Profile:      p,
							Accel:        accel,
							AudioEncoder: ffmpeg.ComponentOptions{Name: "drop"},
							Muxer:        ffmpeg.ComponentOptions{Name: muxer},
						}
						opts = append(opts, o)
					}
					return opts
				}
				out := profs2opts(profiles)
				t := time.Now()
				_, err := tc.Transcode(in, out)
				end := time.Now()
				if err != nil {
					glog.Fatalf("Transcoding failed for session %d segment %d: %v", k, j, err)
				}
				fmt.Printf("%s,%d,%d,%0.4v,%0.4v\n", end.Format("2006-01-02 15:04:05.9999"), k, j, v.Duration, end.Sub(t).Seconds())
				segTxDur := end.Sub(t).Seconds()
				mu.Lock()
				transcodeDur += segTxDur
				srcDur += v.Duration
				segCount++
				if segTxDur <= v.Duration {
					realTimeSegCount += 1
				}
				mu.Unlock()
				iterEnd := time.Now()
				segDur := time.Duration(v.Duration * float64(time.Second))
				if *live {
					time.Sleep(segDur - iterEnd.Sub(iterStart))
				}
			}
			tc.StopTranscoder()
			wg.Done()
		}(i, &wg)
		time.Sleep(300 * time.Millisecond)
	}
	wg.Wait()
	if segCount == 0 || srcDur == 0.0 {
		glog.Fatal("Input manifest has no segments or total duration is 0s")
	}
	statsTable := tablewriter.NewWriter(os.Stderr)
	stats := [][]string{
		{"Concurrent Sessions", fmt.Sprintf("%v", *concurrentSessions)},
		{"Total Segs Transcoded", fmt.Sprintf("%v", segCount)},
		{"Real-Time Segs Transcoded", fmt.Sprintf("%v", realTimeSegCount)},
		{"* Real-Time Segs Ratio *", fmt.Sprintf("%0.4v", float64(realTimeSegCount)/float64(segCount))},
		{"Total Source Duration", fmt.Sprintf("%vs", srcDur)},
		{"Total Transcoding Duration", fmt.Sprintf("%vs", transcodeDur)},
		{"* Real-Time Duration Ratio *", fmt.Sprintf("%0.4v", transcodeDur/srcDur)},
	}

	statsTable.SetAlignment(tablewriter.ALIGN_LEFT)
	statsTable.SetCenterSeparator("*")
	statsTable.SetColumnSeparator("|")
	statsTable.AppendBulk(stats)
	statsTable.Render()
}

func parseVideoProfiles(inp string) []ffmpeg.VideoProfile {
	type profilesJson struct {
		Profiles []struct {
			Name    string `json:"name"`
			Width   int    `json:"width"`
			Height  int    `json:"height"`
			Bitrate int    `json:"bitrate"`
			FPS     uint   `json:"fps"`
			FPSDen  uint   `json:"fpsDen"`
			Profile string `json:"profile"`
			GOP     string `json:"gop"`
		} `json:"profiles"`
	}
	profs := []ffmpeg.VideoProfile{}
	if inp != "" {
		// try opening up json file with profiles
		content, err := ioutil.ReadFile(inp)
		if err == nil && len(content) > 0 {
			// parse json profiles
			resp := &profilesJson{}
			err = json.Unmarshal(content, &resp.Profiles)
			if err != nil {
				glog.Fatal("Unable to unmarshal the passed transcoding option: ", err)
			}
			for _, profile := range resp.Profiles {
				name := profile.Name
				if name == "" {
					name = "custom_" + common.DefaultProfileName(
						profile.Width,
						profile.Height,
						profile.Bitrate)
				}
				var gop time.Duration
				if profile.GOP != "" {
					if profile.GOP == "intra" {
						gop = ffmpeg.GOPIntraOnly
					} else {
						gopFloat, err := strconv.ParseFloat(profile.GOP, 64)
						if err != nil {
							glog.Fatal("Cannot parse the GOP value in the transcoding options: ", err)
						}
						if gopFloat <= 0.0 {
							glog.Fatalf("Invalid gop value %f. Please set it to a positive value", gopFloat)
						}
						gop = time.Duration(gopFloat * float64(time.Second))
					}
				}
				encodingProfile, err := common.EncoderProfileNameToValue(profile.Profile)
				if err != nil {
					glog.Fatal("Unable to parse the H264 encoder profile: ", err)
				}
				prof := ffmpeg.VideoProfile{
					Name:         name,
					Bitrate:      fmt.Sprint(profile.Bitrate),
					Framerate:    profile.FPS,
					FramerateDen: profile.FPSDen,
					Resolution:   fmt.Sprintf("%dx%d", profile.Width, profile.Height),
					Profile:      encodingProfile,
					GOP:          gop,
				}
				profs = append(profs, prof)
			}
		} else {
			// check the built-in profiles
			profs = make([]ffmpeg.VideoProfile, 0)
			presets := strings.Split(inp, ",")
			for _, v := range presets {
				if p, ok := ffmpeg.VideoProfileLookup[strings.TrimSpace(v)]; ok {
					profs = append(profs, p)
				}
			}
		}
		if len(profs) <= 0 {
			glog.Fatalf("No transcoding options provided")
		}
	}
	return profs
}
<commit_msg>cmd/livepeer_bench: Benchmark content detection<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strconv"
	"strings"
	"sync"
	"time"

	//"runtime/pprof"

	"github.com/golang/glog"
	"github.com/livepeer/go-livepeer/common"
	"github.com/livepeer/lpms/ffmpeg"
	"github.com/livepeer/m3u8"
	"github.com/olekukonko/tablewriter"
)

func main() {
	// Override the default flag set since there are dependencies that
	// incorrectly add their own flags (specifically, due to the 'testing'
	// package being linked)
	flag.Set("logtostderr", "true")
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)

	in := flag.String("in", "", "Input m3u8 manifest file")
	live := flag.Bool("live", true, "Simulate live stream")
	concurrentSessions := flag.Int("concurrentSessions", 1, "# of concurrent transcode sessions")
	segs := flag.Int("segs", 0, "Maximum # of segments to transcode (default all)")
	transcodingOptions := flag.String("transcodingOptions", "P240p30fps16x9,P360p30fps16x9,P720p30fps16x9", "Transcoding options for broadcast job, or path to json config")
	nvidia := flag.String("nvidia", "", "Comma-separated list of Nvidia GPU device IDs (or \"all\" for all available devices)")
	outPrefix := flag.String("outPrefix", "", "Output segments' prefix (no segments are generated by default)")
	detectionFreq := flag.Int("detectionFreq", 0, "Run content-detection on every nth segment. No detection occurs for default frequency of 0.")
	detectionSampleRate := flag.Uint("detectionSampleRate", 1, "Run content-detection on every nth frame of a particular segment, if detectionFreq > 0.")

	flag.Parse()

	if *in == "" {
		glog.Errorf("Please provide the input manifest as `%s -in <input.m3u8>`", os.Args[0])
		flag.Usage()
		os.Exit(1)
	}

	profiles := parseVideoProfiles(*transcodingOptions)

	f, err := os.Open(*in)
	if err != nil {
		glog.Fatal("Couldn't open input manifest: ", err)
	}
	p, _, err := m3u8.DecodeFrom(bufio.NewReader(f), true)
	if err != nil {
		glog.Fatal("Couldn't decode input manifest: ", err)
	}
	pl, ok := p.(*m3u8.MediaPlaylist)
	if !ok {
		glog.Fatalf("Expecting media playlist in the input %s", *in)
	}

	accel := ffmpeg.Software
	devices := []string{}
	if *nvidia != "" {
		var err error
		accel = ffmpeg.Nvidia
		devices, err = common.ParseNvidiaDevices(*nvidia)
		if err != nil {
			glog.Fatalf("Error while parsing '-nvidia %v' flag: %v", *nvidia, err)
		}
	}
	var wg sync.WaitGroup
	dir := path.Dir(*in)

	table := tablewriter.NewWriter(os.Stderr)
	data := [][]string{
		{"Source File", *in},
		{"Transcoding Options", *transcodingOptions},
		{"Concurrent Sessions", fmt.Sprintf("%v", *concurrentSessions)},
		{"Live Mode", fmt.Sprintf("%v", *live)},
	}

	if accel == ffmpeg.Nvidia && len(devices) > 0 {
		data = append(data, []string{"Nvidia GPU IDs", fmt.Sprintf("%v", strings.Join(devices, ","))})
	}
	if *detectionFreq > 0 {
		data = append(data, []string{"Content Detection (segment_freq,frame_sample_rate)", fmt.Sprintf("%v,%v", *detectionFreq, *detectionSampleRate)})
	}

	table.SetAlignment(tablewriter.ALIGN_LEFT)
	table.SetCenterSeparator("*")
	table.SetColumnSeparator("|")
	table.AppendBulk(data)
	table.Render()

	detectionOpts := ffmpeg.TranscodeOptions{
		Accel:        accel,
		AudioEncoder: ffmpeg.ComponentOptions{Name: "copy"},
		Detector: &ffmpeg.SceneClassificationProfile{
			SampleRate: *detectionSampleRate,
			Input:      ffmpeg.DSceneAdultSoccer.Input,
			Output:     ffmpeg.DSceneAdultSoccer.Output,
			Classes:    ffmpeg.DSceneAdultSoccer.Classes,
		},
	}
	if *detectionFreq > 0 {
		err = ffmpeg.InitFFmpegWithDetectorProfile(detectionOpts.Detector, *nvidia)
		if err != nil {
			glog.Fatalf("Could not initialize detector profiles")
		}
		defer ffmpeg.ReleaseFFmpegDetectorProfile()
		fmt.Println("timestamp,session,segment,seg_dur,transcode_time,detect_data")
	} else {
		ffmpeg.InitFFmpeg()
		fmt.Println("timestamp,session,segment,seg_dur,transcode_time")
	}

	segCount := 0
	realTimeSegCount := 0
	srcDur := 0.0
	var mu sync.Mutex
	transcodeDur := 0.0
	for i := 0; i < *concurrentSessions; i++ {
		wg.Add(1)
		go func(k int, wg *sync.WaitGroup) {
			tc := ffmpeg.NewTranscoder()
			for j, v := range pl.Segments {
				iterStart := time.Now()
				if *segs > 0 && j >= *segs {
					break
				}
				if v == nil {
					continue
				}
				u := path.Join(dir, v.URI)
				in := &ffmpeg.TranscodeOptionsIn{
					Fname: u,
					Accel: accel,
				}
				if ffmpeg.Software != accel {
					in.Device = devices[k%len(devices)]
				}
				profs2opts := func(profs []ffmpeg.VideoProfile) []ffmpeg.TranscodeOptions {
					opts := []ffmpeg.TranscodeOptions{}
					for n, p := range profs {
						oname := ""
						muxer := ""
						if *outPrefix != "" {
							oname = fmt.Sprintf("%s_%s_%d_%d_%d.ts", *outPrefix, p.Name, n, k, j)
							muxer = "mpegts"
						} else {
							oname = "-"
							muxer = "null"
						}
						o := ffmpeg.TranscodeOptions{
							Oname:        oname,
							Profile:      p,
							Accel:        accel,
							AudioEncoder: ffmpeg.ComponentOptions{Name: "copy"},
							Muxer:        ffmpeg.ComponentOptions{Name: muxer},
						}
						opts = append(opts, o)
					}
					// add detector profile if freq > 0
					if *detectionFreq > 0 && j%*detectionFreq == 0 {
						opts = append(opts, detectionOpts)
					}
					return opts
				}
				out := profs2opts(profiles)
				t := time.Now()
				res, err := tc.Transcode(in, out)
				end := time.Now()
				if err != nil {
					glog.Fatalf("Transcoding failed for session %d segment %d: %v", k, j, err)
				}
				if *detectionFreq > 0 && j%*detectionFreq == 0 {
					fmt.Printf("%s,%d,%d,%0.4v,%0.4v,%v\n", end.Format("2006-01-02 15:04:05.9999"), k, j, v.Duration, end.Sub(t).Seconds(), res.Encoded[len(res.Encoded)-1].DetectData)
				} else {
					fmt.Printf("%s,%d,%d,%0.4v,%0.4v\n", end.Format("2006-01-02 15:04:05.9999"), k, j, v.Duration, end.Sub(t).Seconds())
				}
				segTxDur := end.Sub(t).Seconds()
				mu.Lock()
				transcodeDur += segTxDur
				srcDur += v.Duration
				segCount++
				if segTxDur <= v.Duration {
					realTimeSegCount += 1
				}
				mu.Unlock()
				iterEnd := time.Now()
				segDur := time.Duration(v.Duration * float64(time.Second))
				if *live {
					time.Sleep(segDur - iterEnd.Sub(iterStart))
				}
			}
			tc.StopTranscoder()
			wg.Done()
		}(i, &wg)
		time.Sleep(2300 * time.Millisecond) // wait for at least one segment before moving on to the next session
	}
	wg.Wait()
	if segCount == 0 || srcDur == 0.0 {
		glog.Fatal("Input manifest has no segments or total duration is 0s")
	}
	statsTable := tablewriter.NewWriter(os.Stderr)
	stats := [][]string{
		{"Concurrent Sessions", fmt.Sprintf("%v", *concurrentSessions)},
		{"Total Segs Transcoded", fmt.Sprintf("%v", segCount)},
		{"Real-Time Segs Transcoded", fmt.Sprintf("%v", realTimeSegCount)},
		{"* Real-Time Segs Ratio *", fmt.Sprintf("%0.4v", float64(realTimeSegCount)/float64(segCount))},
		{"Total Source Duration", fmt.Sprintf("%vs", srcDur)},
		{"Total Transcoding Duration", fmt.Sprintf("%vs", transcodeDur)},
		{"* Real-Time Duration Ratio *", fmt.Sprintf("%0.4v", transcodeDur/srcDur)},
	}

	statsTable.SetAlignment(tablewriter.ALIGN_LEFT)
	statsTable.SetCenterSeparator("*")
	statsTable.SetColumnSeparator("|")
	statsTable.AppendBulk(stats)
	statsTable.Render()
}

func parseVideoProfiles(inp string) []ffmpeg.VideoProfile {
	type profilesJson struct {
		Profiles []struct {
			Name    string `json:"name"`
			Width   int    `json:"width"`
			Height  int    `json:"height"`
			Bitrate int    `json:"bitrate"`
			FPS     uint   `json:"fps"`
			FPSDen  uint   `json:"fpsDen"`
			Profile string `json:"profile"`
			GOP     string `json:"gop"`
		} `json:"profiles"`
	}
	profs := []ffmpeg.VideoProfile{}
	if inp != "" {
		// try opening up json file with profiles
		content, err := ioutil.ReadFile(inp)
		if err == nil && len(content) > 0 {
			// parse json profiles
			resp := &profilesJson{}
			err = json.Unmarshal(content, &resp.Profiles)
			if err != nil {
				glog.Fatal("Unable to unmarshal the passed transcoding option: ", err)
			}
			for _, profile := range resp.Profiles {
				name := profile.Name
				if name == "" {
					name = "custom_" + common.DefaultProfileName(
						profile.Width,
						profile.Height,
						profile.Bitrate)
				}
				var gop time.Duration
				if profile.GOP != "" {
					if profile.GOP == "intra" {
						gop = ffmpeg.GOPIntraOnly
					} else {
						gopFloat, err := strconv.ParseFloat(profile.GOP, 64)
						if err != nil {
							glog.Fatal("Cannot parse the GOP value in the transcoding options: ", err)
						}
						if gopFloat <= 0.0 {
							glog.Fatalf("Invalid gop value %f. Please set it to a positive value", gopFloat)
						}
						gop = time.Duration(gopFloat * float64(time.Second))
					}
				}
				encodingProfile, err := common.EncoderProfileNameToValue(profile.Profile)
				if err != nil {
					glog.Fatal("Unable to parse the H264 encoder profile: ", err)
				}
				prof := ffmpeg.VideoProfile{
					Name:         name,
					Bitrate:      fmt.Sprint(profile.Bitrate),
					Framerate:    profile.FPS,
					FramerateDen: profile.FPSDen,
					Resolution:   fmt.Sprintf("%dx%d", profile.Width, profile.Height),
					Profile:      encodingProfile,
					GOP:          gop,
				}
				profs = append(profs, prof)
			}
		} else {
			// check the built-in profiles
			profs = make([]ffmpeg.VideoProfile, 0)
			presets := strings.Split(inp, ",")
			for _, v := range presets {
				if p, ok := ffmpeg.VideoProfileLookup[strings.TrimSpace(v)]; ok {
					profs = append(profs, p)
				}
			}
		}
		if len(profs) <= 0 {
			glog.Fatalf("No transcoding options provided")
		}
	}
	return profs
}
<|endoftext|>


<commit_before>package extract

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"os"
	"reflect"
	"regexp"
	"strings"

	"github.com/PuerkitoBio/goquery"
	scrape "github.com/slotix/dataflowkit/scrape"
	"golang.org/x/net/html"
)

var logger *log.Logger

func init() {
	logger = log.New(os.Stdout, "extractor: ", log.Lshortfile)
}

// Const is a PieceExtractor that returns a constant value.
type Const struct {
	// The value to return when the Extract() function is called.
	Val interface{}
}

func (e Const) Extract(sel *goquery.Selection) (interface{}, error) {
	return e.Val, nil
}

var _ scrape.PieceExtractor = Const{}

// Text is a PieceExtractor that returns the combined text contents of
// the given selection.
type Text struct {
	// If text is empty in the selection, then return the empty string from Extract,
	// instead of 'nil'.  This signals that the result of this Piece
	// should be included to the results, as opposed to omiting the
	// empty string.
	IncludeIfEmpty bool
}

func (e Text) Extract(sel *goquery.Selection) (interface{}, error) {
	if sel.Text() == "" && !e.IncludeIfEmpty {
		return nil, nil
	}
	return sel.Text(), nil
}

func (e *Text) fillParams(m map[string]interface{}) error {
	err := FillStruct(m, e)
	if err != nil {
		return err
	}
	return nil
}

var _ scrape.PieceExtractor = Text{}

// MultipleText is a PieceExtractor that extracts the text from each element
// in the given selection and returns the texts as an array.
type MultipleText struct {
	// If there are no items in the selection, then return empty list from Extract,
	// instead of the 'nil'.  This signals that the result of this Piece
	// should be included to the results, as opposed to omiting the
	// empty list.
	IncludeIfEmpty bool
}

func (e MultipleText) Extract(sel *goquery.Selection) (interface{}, error) {
	results := []string{}

	sel.Each(func(i int, s *goquery.Selection) {
		results = append(results, s.Text())
	})

	if len(results) == 0 && !e.IncludeIfEmpty {
		return nil, nil
	}

	return results, nil
}

// Html extracts and returns the HTML from inside each element of the
// given selection, as a string.
//
// Note that this results in what is effectively the innerHTML of the element -
// i.e. if our selection consists of ["
ONE
", "
TWO
"]
// then the output will be: "ONETWO".
//
// The return type is a string of all the inner HTML joined together.
type Html struct{}

func (e Html) Extract(sel *goquery.Selection) (interface{}, error) {
	var ret, h string
	var err error

	sel.EachWithBreak(func(i int, s *goquery.Selection) bool {
		h, err = s.Html()
		if err != nil {
			return false
		}

		ret += h
		return true
	})

	if err != nil {
		return nil, err
	}
	return ret, nil
}

var _ scrape.PieceExtractor = Html{}

// OuterHtml extracts and returns the HTML of each element of the
// given selection, as a string.
//
// To illustrate, if our selection consists of
// ["
ONE
", "
TWO
"] then the output will be:
// "
ONE

TWO
".
//
// The return type is a string of all the outer HTML joined together.
type OuterHtml struct{}

func (e OuterHtml) Extract(sel *goquery.Selection) (interface{}, error) {
	output := bytes.NewBufferString("")
	for _, node := range sel.Nodes {
		if err := html.Render(output, node); err != nil {
			return nil, err
		}
	}

	return output.String(), nil
}

var _ scrape.PieceExtractor = OuterHtml{}

// Regex runs the given regex over the contents of each element in the
// given selection, and, for each match, extracts the given subexpression.
// The return type of the extractor is a list of string matches (i.e. []string).
type Regex struct {
	// The regular expression to match.  This regular expression must define
	// exactly one parenthesized subexpression (sometimes known as a "capturing
	// group"), which will be extracted.
	Regex *regexp.Regexp
	// The subexpression of the regex to match.  If this value is not set, and if
	// the given regex has more than one subexpression, an error will be thrown.
	Subexpression int

	// When OnlyText is true, only run the given regex over the text contents of
	// each element in the selection, as opposed to the HTML contents.
	OnlyText bool

	// By default, if there is only a single match, Regex will return
	// the match itself (as opposed to an array containing the single match).
	// Set AlwaysReturnList to true to disable this behaviour, ensuring that the
	// Extract function always returns an array.
	AlwaysReturnList bool

	// If no matches of the provided regex could be extracted, then return the empty list
	// from Extract, instead of 'nil'.  This signals that the result of
	// this Piece should be included to the results, as opposed to
	// omiting the empty list.
	IncludeIfEmpty bool
}

func (e Regex) Extract(sel *goquery.Selection) (interface{}, error) {
	if e.Regex == nil {
		return nil, errors.New("no regex given")
	}
	if e.Regex.NumSubexp() == 0 {
		return nil, errors.New("regex has no subexpressions")
	}

	var subexp int
	if e.Subexpression == 0 {
		if e.Regex.NumSubexp() != 1 {
			e := fmt.Errorf(
				"regex has more than one subexpression (%d), but which to "+
					"extract was not specified",
				e.Regex.NumSubexp())
			return nil, e
		}

		subexp = 1
	} else {
		subexp = e.Subexpression
	}

	results := []string{}

	// For each element in the selector...
	var err error
	sel.EachWithBreak(func(i int, s *goquery.Selection) bool {
		var contents string
		if e.OnlyText {
			contents = s.Text()
		} else {
			contents, err = s.Html()
			if err != nil {
				return false
			}
		}

		ret := e.Regex.FindAllStringSubmatch(contents, -1)

		// A return value of nil == no match
		if ret == nil {
			return true
		}

		// For each regex match...
		for _, submatches := range ret {
			// The 0th entry will be the match of the entire string.  The 1st
			// entry will be the first capturing group, which is what we want to
			// extract.
			if len(submatches) > 1 {
				results = append(results, submatches[subexp])
			}
		}

		return true
	})

	if err != nil {
		return nil, err
	}
	if len(results) == 0 && !e.IncludeIfEmpty {
		return nil, nil
	}
	if len(results) == 1 && !e.AlwaysReturnList {
		return results[0], nil
	}

	return results, nil
}

func (e *Regex) fillParams(m map[string]interface{}) error {
	err := FillStruct(m, e)
	if err != nil {
		return err
	}
	regExp := m["regexp"]
	e.Regex = regexp.MustCompile(regExp.(string))
	return nil
}

var _ scrape.PieceExtractor = Regex{}

// Attr extracts the value of a given HTML attribute from each element
// in the selection, and returns them as a list.
// The return type of the extractor is a list of attribute valueus (i.e. []string).
type Attr struct {
	// The HTML attribute to extract from each element.
	Attr string

	// By default, if there is only a single attribute extracted, AttrExtractor
	// will return the match itself (as opposed to an array containing the single
	// match). Set AlwaysReturnList to true to disable this behaviour, ensuring
	// that the Extract function always returns an array.
	AlwaysReturnList bool

	// If no elements with this attribute are found, then return the empty list from
	// Extract, instead of  'nil'.  This signals that the result of this
	// Piece should be included to the results, as opposed to omiting
	// the empty list.
	IncludeIfEmpty bool
}

func (e Attr) Extract(sel *goquery.Selection) (interface{}, error) {
	if len(e.Attr) == 0 {
		return nil, errors.New("no attribute provided")
	}

	results := []string{}

	sel.Each(func(i int, s *goquery.Selection) {
		if val, found := s.Attr(e.Attr); found {
			results = append(results, val)
		}
	})

	if len(results) == 0 && !e.IncludeIfEmpty {
		return nil, nil
	}
	if len(results) == 1 && !e.AlwaysReturnList {
		return results[0], nil
	}

	return results, nil
}

func (e *Attr) fillParams(m map[string]interface{}) error {
	err := FillStruct(m, e)
	if err != nil {
		return err
	}
	return nil
}

var _ scrape.PieceExtractor = Attr{}

// Count extracts the count of elements that are matched and returns it.
type Count struct {
	// If no elements with this attribute are found, then return a number from
	// Extract, instead of 'nil'.  This signals that the result of this
	// Piece should be included to the results, as opposed to omiting
	// the empty list.
	IncludeIfEmpty bool
}

func (e Count) Extract(sel *goquery.Selection) (interface{}, error) {
	l := sel.Length()
	if l == 0 && !e.IncludeIfEmpty {
		return nil, nil
	}

	return l, nil
}

func FillParams(t string, m map[string]interface{}) (scrape.PieceExtractor, error) {
	var err error
	/*
		var e scrape.PieceExtractor
		switch t {
		case "text":
			e = Text{}
		case "attr":
			e = Attr{}
		case "regex":
			e = Regex{}
		}
		if m != nil {
			err := FillStruct(m, &e)
			if err != nil {
				return nil, err
			}
		}
		return e, nil
	*/

	switch t {
	case "text":
		txt := Text{}
		if m != nil {
			err = txt.fillParams(m)
			if err != nil {
				return nil, err
			}
		}
		return txt, nil
	case "attr":
		a := Attr{}
		if m != nil {
			err = a.fillParams(m)
			if err != nil {
				return nil, err
			}
		}
		return a, nil
	case "regex":
		r := Regex{}
		if m != nil {
			err = r.fillParams(m)
			if err != nil {
				return nil, err
			}
		}
		return r, nil
	}
	return nil, err

}

func FillStruct(m map[string]interface{}, s interface{}) error {
	for k, v := range m {
		err := SetField(s, k, v)
		if err != nil {
			return err
		}
	}
	return nil
}

func SetField(obj interface{}, name string, value interface{}) error {
	structValue := reflect.ValueOf(obj).Elem()
	//structFieldValue := structValue.FieldByName(name)
	structFieldValue := structValue.FieldByName(strings.Title(name))

	if !structFieldValue.IsValid() {
		//skip non-existent fields
		return nil
		//return fmt.Errorf("No such field: %s in obj", name)
	}

	if !structFieldValue.CanSet() {
		return fmt.Errorf("Cannot set %s field value", name)
	}

	structFieldType := structFieldValue.Type()
	val := reflect.ValueOf(value)
	if structFieldType != val.Type() {
		invalidTypeError := errors.New("Provided value type didn't match obj field type")
		return invalidTypeError
	}

	structFieldValue.Set(val)
	return nil
}
<commit_msg>extractors<commit_after>package extract

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"os"
	"reflect"
	"regexp"
	"strings"

	"github.com/PuerkitoBio/goquery"
	scrape "github.com/slotix/dataflowkit/scrape"
	"golang.org/x/net/html"
)

var logger *log.Logger

func init() {
	logger = log.New(os.Stdout, "extractor: ", log.Lshortfile)
}

// Const is a PieceExtractor that returns a constant value.
type Const struct {
	// The value to return when the Extract() function is called.
	Val interface{}
}

func (e Const) Extract(sel *goquery.Selection) (interface{}, error) {
	return e.Val, nil
}

var _ scrape.PieceExtractor = Const{}

// Text is a PieceExtractor that returns the combined text contents of
// the given selection.
type Text struct {
	// If text is empty in the selection, then return the empty string from Extract,
	// instead of 'nil'.  This signals that the result of this Piece
	// should be included to the results, as opposed to omiting the
	// empty string.
	IncludeIfEmpty bool
}

func (e Text) Extract(sel *goquery.Selection) (interface{}, error) {
	if sel.Text() == "" && !e.IncludeIfEmpty {
		return nil, nil
	}
	return sel.Text(), nil
}

func (e *Text) fillParams(m map[string]interface{}) error {
	err := FillStruct(m, e)
	if err != nil {
		return err
	}
	return nil
}

var _ scrape.PieceExtractor = Text{}

// MultipleText is a PieceExtractor that extracts the text from each element
// in the given selection and returns the texts as an array.
type MultipleText struct {
	// If there are no items in the selection, then return empty list from Extract,
	// instead of the 'nil'.  This signals that the result of this Piece
	// should be included to the results, as opposed to omiting the
	// empty list.
	IncludeIfEmpty bool
}

func (e MultipleText) Extract(sel *goquery.Selection) (interface{}, error) {
	results := []string{}

	sel.Each(func(i int, s *goquery.Selection) {
		results = append(results, s.Text())
	})

	if len(results) == 0 && !e.IncludeIfEmpty {
		return nil, nil
	}

	return results, nil
}

// Html extracts and returns the HTML from inside each element of the
// given selection, as a string.
//
// Note that this results in what is effectively the innerHTML of the element -
// i.e. if our selection consists of ["
ONE
", "
TWO
"]
// then the output will be: "ONETWO".
//
// The return type is a string of all the inner HTML joined together.
type Html struct{}

func (e Html) Extract(sel *goquery.Selection) (interface{}, error) {
	var ret, h string
	var err error

	sel.EachWithBreak(func(i int, s *goquery.Selection) bool {
		h, err = s.Html()
		if err != nil {
			return false
		}

		ret += h
		return true
	})

	if err != nil {
		return nil, err
	}
	return ret, nil
}

var _ scrape.PieceExtractor = Html{}

// OuterHtml extracts and returns the HTML of each element of the
// given selection, as a string.
//
// To illustrate, if our selection consists of
// ["
ONE
", "
TWO
"] then the output will be:
// "
ONE

TWO
".
//
// The return type is a string of all the outer HTML joined together.
type OuterHtml struct{}

func (e OuterHtml) Extract(sel *goquery.Selection) (interface{}, error) {
	output := bytes.NewBufferString("")
	for _, node := range sel.Nodes {
		if err := html.Render(output, node); err != nil {
			return nil, err
		}
	}

	return output.String(), nil
}

var _ scrape.PieceExtractor = OuterHtml{}

// Regex runs the given regex over the contents of each element in the
// given selection, and, for each match, extracts the given subexpression.
// The return type of the extractor is a list of string matches (i.e. []string).
type Regex struct {
	// The regular expression to match.  This regular expression must define
	// exactly one parenthesized subexpression (sometimes known as a "capturing
	// group"), which will be extracted.
	Regex *regexp.Regexp
	// The subexpression of the regex to match.  If this value is not set, and if
	// the given regex has more than one subexpression, an error will be thrown.
	Subexpression int

	// When OnlyText is true, only run the given regex over the text contents of
	// each element in the selection, as opposed to the HTML contents.
	OnlyText bool

	// By default, if there is only a single match, Regex will return
	// the match itself (as opposed to an array containing the single match).
	// Set AlwaysReturnList to true to disable this behaviour, ensuring that the
	// Extract function always returns an array.
	AlwaysReturnList bool

	// If no matches of the provided regex could be extracted, then return the empty list
	// from Extract, instead of 'nil'.  This signals that the result of
	// this Piece should be included to the results, as opposed to
	// omiting the empty list.
	IncludeIfEmpty bool
}

func (e Regex) Extract(sel *goquery.Selection) (interface{}, error) {
	if e.Regex == nil {
		return nil, errors.New("no regex given")
	}
	if e.Regex.NumSubexp() == 0 {
		return nil, errors.New("regex has no subexpressions")
	}

	var subexp int
	if e.Subexpression == 0 {
		if e.Regex.NumSubexp() != 1 {
			e := fmt.Errorf(
				"regex has more than one subexpression (%d), but which to "+
					"extract was not specified",
				e.Regex.NumSubexp())
			return nil, e
		}

		subexp = 1
	} else {
		subexp = e.Subexpression
	}

	results := []string{}

	// For each element in the selector...
	var err error
	sel.EachWithBreak(func(i int, s *goquery.Selection) bool {
		var contents string
		if e.OnlyText {
			contents = s.Text()
		} else {
			contents, err = s.Html()
			if err != nil {
				return false
			}
		}

		ret := e.Regex.FindAllStringSubmatch(contents, -1)

		// A return value of nil == no match
		if ret == nil {
			return true
		}

		// For each regex match...
		for _, submatches := range ret {
			// The 0th entry will be the match of the entire string.  The 1st
			// entry will be the first capturing group, which is what we want to
			// extract.
			if len(submatches) > 1 {
				results = append(results, submatches[subexp])
			}
		}

		return true
	})

	if err != nil {
		return nil, err
	}
	if len(results) == 0 && !e.IncludeIfEmpty {
		return nil, nil
	}
	if len(results) == 1 && !e.AlwaysReturnList {
		return results[0], nil
	}

	return results, nil
}

func (e *Regex) fillParams(m map[string]interface{}) error {
	err := FillStruct(m, e)
	if err != nil {
		return err
	}
	regExp := m["regexp"]
	e.Regex = regexp.MustCompile(regExp.(string))
	return nil
}

var _ scrape.PieceExtractor = Regex{}

// Attr extracts the value of a given HTML attribute from each element
// in the selection, and returns them as a list.
// The return type of the extractor is a list of attribute valueus (i.e. []string).
type Attr struct {
	// The HTML attribute to extract from each element.
	Attr string

	// By default, if there is only a single attribute extracted, AttrExtractor
	// will return the match itself (as opposed to an array containing the single
	// match). Set AlwaysReturnList to true to disable this behaviour, ensuring
	// that the Extract function always returns an array.
	AlwaysReturnList bool

	// If no elements with this attribute are found, then return the empty list from
	// Extract, instead of  'nil'.  This signals that the result of this
	// Piece should be included to the results, as opposed to omiting
	// the empty list.
	IncludeIfEmpty bool
}

func (e Attr) Extract(sel *goquery.Selection) (interface{}, error) {
	if len(e.Attr) == 0 {
		return nil, errors.New("no attribute provided")
	}

	results := []string{}

	sel.Each(func(i int, s *goquery.Selection) {
		if val, found := s.Attr(e.Attr); found {
			results = append(results, val)
		}
	})

	if len(results) == 0 && !e.IncludeIfEmpty {
		return nil, nil
	}
	if len(results) == 1 && !e.AlwaysReturnList {
		return results[0], nil
	}

	return results, nil
}

func (e *Attr) fillParams(m map[string]interface{}) error {
	err := FillStruct(m, e)
	if err != nil {
		return err
	}
	return nil
}

var _ scrape.PieceExtractor = Attr{}

// Count extracts the count of elements that are matched and returns it.
type Count struct {
	// If no elements with this attribute are found, then return a number from
	// Extract, instead of 'nil'.  This signals that the result of this
	// Piece should be included to the results, as opposed to omiting
	// the empty list.
	IncludeIfEmpty bool
}

func (e Count) Extract(sel *goquery.Selection) (interface{}, error) {
	l := sel.Length()
	if l == 0 && !e.IncludeIfEmpty {
		return nil, nil
	}

	return l, nil
}

func FillParams(t string, m map[string]interface{}) (scrape.PieceExtractor, error) {
	//var err error

	logger.Println(t)
	var e scrape.PieceExtractor
	switch t {
	case "text":
		e = &Text{}
	case "attr":
		e = &Attr{}
	case "regex":
		//e = &Regex{}
		r := &Regex{}
		regExp := m["regexp"]
		r.Regex = regexp.MustCompile(regExp.(string))
		e = r
	}
	if m != nil {
		err := FillStruct(m, e)
		if err != nil {
			return nil, err
		}
	}
	return e, nil

	/*
		switch t {
		case "text":
			txt := Text{}
			if m != nil {
				err = txt.fillParams(m)
				if err != nil {
					return nil, err
				}
			}
			return txt, nil
		case "attr":
			a := Attr{}
			if m != nil {
				err = a.fillParams(m)
				if err != nil {
					return nil, err
				}
			}
			return a, nil
		case "regex":
			r := Regex{}
			if m != nil {
				err = r.fillParams(m)
				if err != nil {
					return nil, err
				}
			}
			return r, nil
		}

		return nil, err
*/
}

func FillStruct(m map[string]interface{}, s interface{}) error {
	for k, v := range m {
		logger.Println(k,v)
		err := SetField(s, k, v)
		if err != nil {
			return err
		}
	}
	return nil
}

func SetField(obj interface{}, name string, value interface{}) error {
//logger.Printf("%T, %t", obj, obj)
	structValue := reflect.ValueOf(obj).Elem()
	//structFieldValue := structValue.FieldByName(name)
	structFieldValue := structValue.FieldByName(strings.Title(name))

	if !structFieldValue.IsValid() {
		//skip non-existent fields
		return nil
		//return fmt.Errorf("No such field: %s in obj", name)
	}

	if !structFieldValue.CanSet() {
		return fmt.Errorf("Cannot set %s field value", name)
	}

	structFieldType := structFieldValue.Type()
	val := reflect.ValueOf(value)
	if structFieldType != val.Type() {
		invalidTypeError := errors.New("Provided value type didn't match obj field type")
		return invalidTypeError
	}

	structFieldValue.Set(val)
	return nil
}
<|endoftext|>