<commit_before>package service

import "fmt"

//Builder is a simple implementation of ContainerBuilder
type Builder struct {
	definitions []Definition
}

//Insert a new definition into the Builder
func (b *Builder) Insert(def Definition) {
	b.definitions = append(b.definitions, def)
}

//Build builds the container once all definitions have been place in it
//dependencies need to be inserted in order for now
func (b *Builder) Build() (Container, error) {
	numDefs := len(b.definitions)
	servs := make(map[string]interface{}, numDefs)
	for _, def := range b.definitions {
		numDeps := len(def.Dependencies)
		deps := make(map[string]interface{}, numDeps)
		for _, name := range def.Dependencies {
			dep, ok := servs[name]
			if !ok {
				return nil, fmt.Errorf("service: Could not find"+
					"dependency %q for service %q. Please make"+
					"sure to insert them in order", name, def.Name)
			}
			deps[name] = dep
		}
		service, err := def.Initializer.Init(deps, def.Configuration)
		if err != nil {
			return nil, err
		}
		servs[def.Name] = service
	}

	return SimpleContainer{
		services: servs,
	}, nil
}
<commit_msg>Fix spaces in error message<commit_after>package service

import "fmt"

//Builder is a simple implementation of ContainerBuilder
type Builder struct {
	definitions []Definition
}

//Insert a new definition into the Builder
func (b *Builder) Insert(def Definition) {
	b.definitions = append(b.definitions, def)
}

//Build builds the container once all definitions have been place in it
//dependencies need to be inserted in order for now
func (b *Builder) Build() (Container, error) {
	numDefs := len(b.definitions)
	servs := make(map[string]interface{}, numDefs)
	for _, def := range b.definitions {
		numDeps := len(def.Dependencies)
		deps := make(map[string]interface{}, numDeps)
		for _, name := range def.Dependencies {
			dep, ok := servs[name]
			if !ok {
				return nil, fmt.Errorf("service: Could not find "+
					"dependency %q for service %q. Please make "+
					"sure to insert them in order", name, def.Name)
			}
			deps[name] = dep
		}
		service, err := def.Initializer.Init(deps, def.Configuration)
		if err != nil {
			return nil, err
		}
		servs[def.Name] = service
	}

	return SimpleContainer{
		services: servs,
	}, nil
}
<|endoftext|>


<commit_before>//Package service contains the logic to build a HTTP/Rest Service for the MiSSy runtime environment
package service

import (
	"bytes"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"syscall"
	"time"

	gctx "github.com/gorilla/context"
	"github.com/gorilla/mux"
	"github.com/microdevs/missy/config"
	"github.com/microdevs/missy/log"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

// Server is implemented by *http.Server
type Server interface {
	ListenAndServe() error
	Shutdowner
}

// TLSServer is implemented by *http.Server
type TLSServer interface {
	ListenAndServeTLS(string, string) error
	Shutdowner
}

// Shutdowner is implemented by *http.Server, and optionally by *http.Server.Handler
type Shutdowner interface {
	Shutdown(ctx context.Context) error
}

// signals is the channel used to signal shutdown
var signals chan os.Signal

type key int

// PrometheusInstance is a key for the context value of the PrometheusHolder
const PrometheusInstance key = 0

// RouterInstance is the key for the goilla/mux router instance in the context
const RouterInstance key = 1

// RequestTimer is the key for the request timer in context
const RequestTimer key = 2

// Service type provides a HTTP/Rest service
type Service struct {
	name       string
	Host       string
	Port       string
	Prometheus *PrometheusHolder
	timer      *Timer
	Router     *mux.Router
	Stop       chan os.Signal
	ServeMux   *http.ServeMux
}

var listenPort = "8080"
var listenHost = "localhost"
var controllerAddr string

// FlagMissyControllerAddressDefault is a default for the missy-controller url used in the during service initialisation when given the init flag
const FlagMissyControllerAddressDefault = "http://missy-controller"

// FlagMissyControllerUsage is a usage message for the missy-controller url used in the during service initialisation when given the init flag
const FlagMissyControllerUsage = "The address of the MiSSy controller"

// init checks for init flag and executes the service registration with the missy controller if applicable
func init() {
	initCmd := flag.NewFlagSet("init", flag.ExitOnError)
	initCmd.StringVar(&controllerAddr, "addr", FlagMissyControllerAddressDefault, FlagMissyControllerUsage)
	initCmd.StringVar(&controllerAddr, "a", FlagMissyControllerAddressDefault, FlagMissyControllerUsage+" (Shorthand)")

	if len(os.Args) > 1 && os.Args[1] == "init" {
		initCmd.Parse(os.Args[2:])
		c := config.GetInstance()
		cjson, jsonErr := json.Marshal(c)
		if jsonErr != nil {
			fmt.Println("Error marshalling config to json.")
			os.Exit(1)
		}
		log.Infof("Registering service %s with MiSSy controller at %s", c.Name, controllerAddr)
		_, err := http.Post(controllerAddr+"/registerService", "application/json", bytes.NewReader(cjson))
		// todo: check response for return status
		if err != nil {
			fmt.Printf("Can not reach missy controller: %s", err)
			os.Exit(1)
		}
		os.Exit(0)
	}

}

// New returns a new Service object
func New() *Service {

	if _, present := os.LookupEnv("LISTEN_HOST"); present {
		listenHost = os.Getenv("LISTEN_HOST")
	}

	if _, present := os.LookupEnv("LISTEN_PORT"); present {
		listenPort = os.Getenv("LISTEN_PORT")
	}

	c := config.GetInstance()

	s := &Service{
		name:       c.Name,
		Host:       listenHost,
		Port:       listenPort,
		Prometheus: NewPrometheus(c.Name),
		Router:     mux.NewRouter(),
		ServeMux:   http.NewServeMux(),
	}
	s.prepareBeforeStart()
	return s
}

// Start starts the http server
func (s *Service) Start() {
	// start the server
	log.Infof("Starting service %s listening on %s:%s ...", s.name, s.Host, s.Port)
	// set host and port to listen to
	listen := s.Host + ":" + s.Port
	h := &http.Server{Addr: listen, Handler: s.Router}
	// run server in background
	go func() {
		err := h.ListenAndServe()
		if err != nil {
			log.Fatalf("Error starting Service due to %v", err)
		}
		log.Debugf("server shut down")
	}()

	s.prepareShutdown(h)
}

func (s *Service) prepareShutdown(h Server) {
	s.Stop = make(chan os.Signal, 1)
	signal.Notify(s.Stop, os.Interrupt, syscall.SIGTERM)
	<-s.Stop
	shutdown(h)
}

// Shutdown allows to stop the HTTP Server gracefully
func (s *Service) Shutdown() {
	s.Stop <- os.Signal(os.Interrupt)
}

func shutdown(s Shutdowner) {
	if s == nil {
		return
	}

	// todo: make configurable
	timeout := time.Second * 5

	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	log.Printf(""+
		"Server shutdown with timeout: %s", timeout)

	if err := s.Shutdown(ctx); err != nil {
		log.Printf("Error: %v", err)
	} else {
		if hs, ok := s.(*http.Server); ok {
			log.Printf("Finished all in-flight HTTP requests")

			if hss, ok := hs.Handler.(Shutdowner); ok {
				select {
				case <-ctx.Done():
					if err := ctx.Err(); err != nil {
						log.Printf("Error: %v", err)
						return
					}
				default:
					if deadline, ok := ctx.Deadline(); ok {
						secs := (time.Until(deadline) + time.Second/2) / time.Second
						log.Printf("Shutting down handler with timeout: %ds", secs)
					}

					done := make(chan error)

					go func() {
						<-ctx.Done()
						done <- ctx.Err()
					}()

					go func() {
						done <- hss.Shutdown(ctx)
					}()

					if err := <-done; err != nil {
						log.Printf("Error: %v", err)
						return
					}
				}
			}
		}

		if deadline, ok := ctx.Deadline(); ok {
			secs := (time.Until(deadline) + time.Second/2) / time.Second
			log.Printf("Shutdown finished %ds before deadline", secs)
		}
	}
}

// prepareBeforeStart sets up the standard handlers
func (s *Service) prepareBeforeStart() {
	s.timer = NewTimer()
	s.Router.Handle("/metrics", promhttp.Handler()).Methods("GET")
	s.Router.HandleFunc("/health", s.healthHandler).Methods("GET")
	s.Router.HandleFunc("/info", s.infoHandler).Methods("GET")

	s.ServeMux.Handle("/", s.Router)

}

// HandleFunc excepts a HanderFunc an converts it to a handler, then registers this handler
// Deprecated: Developers should use SecureHandleFunc() or UnsafeHandleFunc() explicitly
func (s *Service) HandleFunc(pattern string, handler func(http.ResponseWriter, *http.Request)) *mux.Route {
	return s.UnsafeHandle(pattern, http.HandlerFunc(handler))
}

// UnsafeHandleFunc excepts a HanderFunc an converts it to a handler, then registers this handler
func (s *Service) UnsafeHandleFunc(pattern string, handler func(http.ResponseWriter, *http.Request)) *mux.Route {
	return s.UnsafeHandle(pattern, http.HandlerFunc(handler))
}

// SecureHandleFunc excepts a HanderFunc an converts it to a handler, then registers this handler
func (s *Service) SecureHandleFunc(pattern string, handler func(http.ResponseWriter, *http.Request)) *mux.Route {
	return s.SecureHandle(pattern, http.HandlerFunc(handler))
}

// Handle is a wrapper around the original Go handle func with logging recovery and metrics
// Deprecated: Developers should use SecureHandle() or UnsafeHandle() explicitly
func (s *Service) Handle(pattern string, originalHandler http.Handler) *mux.Route {
	h := s.makeHandler(originalHandler, false)
	return s.Router.Handle(pattern, h)
}

// UnsafeHandle is a wrapper around the original Go handle func with logging recovery and metrics
func (s *Service) UnsafeHandle(pattern string, originalHandler http.Handler) *mux.Route {
	h := s.makeHandler(originalHandler, false)
	return s.Router.Handle(pattern, h)
}

// SecureHandle is a wrapper around the original Go handle func with logging recovery and metrics
func (s *Service) SecureHandle(pattern string, originalHandler http.Handler) *mux.Route {
	initPublicKey()
	h := s.makeHandler(originalHandler, true)
	return s.Router.Handle(pattern, h)
}

// Makes a handler that wraps Missy specific functionality and returns either a secure or insecure chain
// a secure chain includes the auth handler
func (s *Service) makeHandler(originalHandler http.Handler, secure bool) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if err := recover(); err != nil {
				stack := make([]byte, 1024*8)
				stack = stack[:runtime.Stack(stack, false)]
				log.Errorf("PANIC: %s\n%s", err, stack)
				http.Error(w, "500 Internal Server Error", http.StatusInternalServerError)
			}
		}()
		// build context
		gctx.Set(r, PrometheusInstance, s.Prometheus)
		gctx.Set(r, RouterInstance, s.Router)
		// call custom handler
		chain := NewChain(StartTimerHandler).Final(FinalHandler).Then(originalHandler)
		if secure {
			chain = NewChain(StartTimerHandler, AuthHandler).Final(FinalHandler).Then(originalHandler)
		}
		chain.ServeHTTP(w, r)
	})
}

// ResponseWriter is the MiSSy owned response writer object
type ResponseWriter struct {
	http.ResponseWriter
	Status int
}

// WriteHeader overrides the original WriteHeader function to keep the status code
func (w *ResponseWriter) WriteHeader(code int) {
	w.Status = code
	w.ResponseWriter.WriteHeader(code)
}

// ResponseWriter wrapper for http.ResponseWriter interface
func (w *ResponseWriter) Write(b []byte) (int, error) {
	return w.ResponseWriter.Write(b)
}

// Header wrapper for http.ResponseWriter interface
func (w *ResponseWriter) Header() http.Header {
	return w.ResponseWriter.Header()
}
<commit_msg>refactor(fmt): formatting<commit_after>//Package service contains the logic to build a HTTP/Rest Service for the MiSSy runtime environment
package service

import (
	"bytes"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"syscall"
	"time"

	gctx "github.com/gorilla/context"
	"github.com/gorilla/mux"
	"github.com/microdevs/missy/config"
	"github.com/microdevs/missy/log"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

// Server is implemented by *http.Server
type Server interface {
	ListenAndServe() error
	Shutdowner
}

// TLSServer is implemented by *http.Server
type TLSServer interface {
	ListenAndServeTLS(string, string) error
	Shutdowner
}

// Shutdowner is implemented by *http.Server, and optionally by *http.Server.Handler
type Shutdowner interface {
	Shutdown(ctx context.Context) error
}

// signals is the channel used to signal shutdown
var signals chan os.Signal

type key int

// PrometheusInstance is a key for the context value of the PrometheusHolder
const PrometheusInstance key = 0

// RouterInstance is the key for the goilla/mux router instance in the context
const RouterInstance key = 1

// RequestTimer is the key for the request timer in context
const RequestTimer key = 2

// Service type provides a HTTP/Rest service
type Service struct {
	name       string
	Host       string
	Port       string
	Prometheus *PrometheusHolder
	timer      *Timer
	Router     *mux.Router
	Stop       chan os.Signal
	ServeMux   *http.ServeMux
}

var listenPort = "8080"
var listenHost = "localhost"
var controllerAddr string

// FlagMissyControllerAddressDefault is a default for the missy-controller url used in the during service initialisation when given the init flag
const FlagMissyControllerAddressDefault = "http://missy-controller"

// FlagMissyControllerUsage is a usage message for the missy-controller url used in the during service initialisation when given the init flag
const FlagMissyControllerUsage = "The address of the MiSSy controller"

// init checks for init flag and executes the service registration with the missy controller if applicable
func init() {
	initCmd := flag.NewFlagSet("init", flag.ExitOnError)
	initCmd.StringVar(&controllerAddr, "addr", FlagMissyControllerAddressDefault, FlagMissyControllerUsage)
	initCmd.StringVar(&controllerAddr, "a", FlagMissyControllerAddressDefault, FlagMissyControllerUsage+" (Shorthand)")

	if len(os.Args) > 1 && os.Args[1] == "init" {
		initCmd.Parse(os.Args[2:])
		c := config.GetInstance()
		cjson, jsonErr := json.Marshal(c)
		if jsonErr != nil {
			fmt.Println("Error marshalling config to json.")
			os.Exit(1)
		}
		log.Infof("Registering service %s with MiSSy controller at %s", c.Name, controllerAddr)
		_, err := http.Post(controllerAddr+"/registerService", "application/json", bytes.NewReader(cjson))
		// todo: check response for return status
		if err != nil {
			fmt.Printf("Can not reach missy controller: %s", err)
			os.Exit(1)
		}
		os.Exit(0)
	}

}

// New returns a new Service object
func New() *Service {

	if _, present := os.LookupEnv("LISTEN_HOST"); present {
		listenHost = os.Getenv("LISTEN_HOST")
	}

	if _, present := os.LookupEnv("LISTEN_PORT"); present {
		listenPort = os.Getenv("LISTEN_PORT")
	}

	c := config.GetInstance()

	s := &Service{
		name:       c.Name,
		Host:       listenHost,
		Port:       listenPort,
		Prometheus: NewPrometheus(c.Name),
		Router:     mux.NewRouter(),
		ServeMux:   http.NewServeMux(),
	}
	s.prepareBeforeStart()
	return s
}

// Start starts the http server
func (s *Service) Start() {
	// start the server
	log.Infof("Starting service %s listening on %s:%s ...", s.name, s.Host, s.Port)
	// set host and port to listen to
	listen := s.Host + ":" + s.Port
	h := &http.Server{Addr: listen, Handler: s.Router}
	// run server in background
	go func() {
		err := h.ListenAndServe()
		if err != nil {
			log.Fatalf("Error starting Service due to %v", err)
		}
		log.Debugf("server shut down")
	}()

	s.prepareShutdown(h)
}

func (s *Service) prepareShutdown(h Server) {
	s.Stop = make(chan os.Signal, 1)
	signal.Notify(s.Stop, os.Interrupt, syscall.SIGTERM)
	<-s.Stop
	shutdown(h)
}

// Shutdown allows to stop the HTTP Server gracefully
func (s *Service) Shutdown() {
	s.Stop <- os.Signal(os.Interrupt)
}

func shutdown(s Shutdowner) {
	if s == nil {
		return
	}

	// todo: make configurable
	timeout := time.Second * 5

	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	log.Printf(""+
		"Server shutdown with timeout: %s", timeout)

	if err := s.Shutdown(ctx); err != nil {
		log.Printf("Error: %v", err)
	} else {
		if hs, ok := s.(*http.Server); ok {
			log.Printf("Finished all in-flight HTTP requests")

			if hss, ok := hs.Handler.(Shutdowner); ok {
				select {
				case <-ctx.Done():
					if err := ctx.Err(); err != nil {
						log.Printf("Error: %v", err)
						return
					}
				default:
					if deadline, ok := ctx.Deadline(); ok {
						secs := (time.Until(deadline) + time.Second/2) / time.Second
						log.Printf("Shutting down handler with timeout: %ds", secs)
					}

					done := make(chan error)

					go func() {
						<-ctx.Done()
						done <- ctx.Err()
					}()

					go func() {
						done <- hss.Shutdown(ctx)
					}()

					if err := <-done; err != nil {
						log.Printf("Error: %v", err)
						return
					}
				}
			}
		}

		if deadline, ok := ctx.Deadline(); ok {
			secs := (time.Until(deadline) + time.Second/2) / time.Second
			log.Printf("Shutdown finished %ds before deadline", secs)
		}
	}
}

// prepareBeforeStart sets up the standard handlers
func (s *Service) prepareBeforeStart() {
	s.timer = NewTimer()
	s.Router.Handle("/metrics", promhttp.Handler()).Methods("GET")
	s.Router.HandleFunc("/health", s.healthHandler).Methods("GET")
	s.Router.HandleFunc("/info", s.infoHandler).Methods("GET")

	s.ServeMux.Handle("/", s.Router)

}

// HandleFunc excepts a HanderFunc an converts it to a handler, then registers this handler
// Deprecated: Developers should use SecureHandleFunc() or UnsafeHandleFunc() explicitly
func (s *Service) HandleFunc(pattern string, handler func(http.ResponseWriter, *http.Request)) *mux.Route {
	return s.UnsafeHandle(pattern, http.HandlerFunc(handler))
}

// UnsafeHandleFunc excepts a HanderFunc an converts it to a handler, then registers this handler
func (s *Service) UnsafeHandleFunc(pattern string, handler func(http.ResponseWriter, *http.Request)) *mux.Route {
	return s.UnsafeHandle(pattern, http.HandlerFunc(handler))
}

// SecureHandleFunc excepts a HanderFunc an converts it to a handler, then registers this handler
func (s *Service) SecureHandleFunc(pattern string, handler func(http.ResponseWriter, *http.Request)) *mux.Route {
	return s.SecureHandle(pattern, http.HandlerFunc(handler))
}

// Handle is a wrapper around the original Go handle func with logging recovery and metrics
// Deprecated: Developers should use SecureHandle() or UnsafeHandle() explicitly
func (s *Service) Handle(pattern string, originalHandler http.Handler) *mux.Route {
	h := s.makeHandler(originalHandler, false)
	return s.Router.Handle(pattern, h)
}

// UnsafeHandle is a wrapper around the original Go handle func with logging recovery and metrics
func (s *Service) UnsafeHandle(pattern string, originalHandler http.Handler) *mux.Route {
	h := s.makeHandler(originalHandler, false)
	return s.Router.Handle(pattern, h)
}

// SecureHandle is a wrapper around the original Go handle func with logging recovery and metrics
func (s *Service) SecureHandle(pattern string, originalHandler http.Handler) *mux.Route {
	initPublicKey()
	h := s.makeHandler(originalHandler, true)
	return s.Router.Handle(pattern, h)
}

// Makes a handler that wraps Missy specific functionality and returns either a secure or insecure chain
// a secure chain includes the auth handler
func (s *Service) makeHandler(originalHandler http.Handler, secure bool) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if err := recover(); err != nil {
				stack := make([]byte, 1024*8)
				stack = stack[:runtime.Stack(stack, false)]
				log.Errorf("PANIC: %s\n%s", err, stack)
				http.Error(w, "500 Internal Server Error", http.StatusInternalServerError)
			}
		}()
		// build context
		gctx.Set(r, PrometheusInstance, s.Prometheus)
		gctx.Set(r, RouterInstance, s.Router)
		// call custom handler
		chain := NewChain(StartTimerHandler).Final(FinalHandler).Then(originalHandler)
		if secure {
			chain = NewChain(StartTimerHandler, AuthHandler).Final(FinalHandler).Then(originalHandler)
		}
		chain.ServeHTTP(w, r)
	})
}

// ResponseWriter is the MiSSy owned response writer object
type ResponseWriter struct {
	http.ResponseWriter
	Status int
}

// WriteHeader overrides the original WriteHeader function to keep the status code
func (w *ResponseWriter) WriteHeader(code int) {
	w.Status = code
	w.ResponseWriter.WriteHeader(code)
}

// ResponseWriter wrapper for http.ResponseWriter interface
func (w *ResponseWriter) Write(b []byte) (int, error) {
	return w.ResponseWriter.Write(b)
}

// Header wrapper for http.ResponseWriter interface
func (w *ResponseWriter) Header() http.Header {
	return w.ResponseWriter.Header()
}
<|endoftext|>


<commit_before>// +build ignore

package main

import (
	"log"
	"net"
)

func main() {
	ifs, err := net.Interfaces()
	if err != nil {
		panic(err)
	}
	for _, if_ := range ifs {
		log.Println(if_)
		addrs, err := if_.Addrs()
		if err != nil {
			panic(err)
		}
		for _, addr := range addrs {
			log.Printf("\t%s %s\n", addr.Network(), addr)
		}
		mcastAddrs, err := if_.MulticastAddrs()
		if err != nil {
			panic(err)
		}
		for _, addr := range mcastAddrs {
			log.Printf("\t%s\n", addr)
		}
	}
}
<commit_msg>play/print-ifs: Print interface info to stdout<commit_after>// +build ignore

package main

import (
	"fmt"
	"net"
)

func main() {
	ifs, err := net.Interfaces()
	if err != nil {
		panic(err)
	}
	for _, if_ := range ifs {
		fmt.Printf("%#v\n", if_)
		addrs, err := if_.Addrs()
		if err != nil {
			panic(err)
		}
		for _, addr := range addrs {
			fmt.Printf("\t%s %s\n", addr.Network(), addr)
		}
		mcastAddrs, err := if_.MulticastAddrs()
		if err != nil {
			panic(err)
		}
		for _, addr := range mcastAddrs {
			fmt.Printf("\t%s\n", addr)
		}
	}
}
<|endoftext|>


<commit_before>package zmq

import (
	"encoding/json"
	"fmt"

	log "github.com/Sirupsen/logrus"
	odplugin "github.com/oursky/skygear/plugin"
	"github.com/oursky/skygear/plugin/common"
	"github.com/oursky/skygear/skyconv"
	"github.com/oursky/skygear/skydb"
	"github.com/zeromq/goczmq"
	"golang.org/x/net/context"
)

const initRequestTimeout = 2000

type zmqTransport struct {
	state       odplugin.TransportState
	name        string
	iaddr       string // the internal addr used by goroutines to make request to plugin
	eaddr       string // the addr exposed for plugin to connect to with REP.
	broker      *Broker
	initHandler odplugin.TransportInitHandler
	logger      *log.Entry
}

type request struct {
	Context context.Context
	Kind    string
	Name    string
	Param   interface{}
	Timeout int // timeout in millisecond
}

type hookRequest struct {
	Record   interface{} `json:"record"`
	Original interface{} `json:"original"`
}

// type-safe constructors for request.Param assignment

func newLambdaRequest(ctx context.Context, name string, args json.RawMessage) *request {
	return &request{Kind: "op", Name: name, Param: args, Context: ctx}
}

func newHandlerRequest(name string, input json.RawMessage) *request {
	return &request{Kind: "handler", Name: name, Param: input}
}

func newHookRequest(trigger string, record *skydb.Record, originalRecord *skydb.Record, ctx context.Context) *request {
	param := hookRequest{
		Record:   (*skyconv.JSONRecord)(record),
		Original: (*skyconv.JSONRecord)(originalRecord),
	}
	return &request{Kind: "hook", Name: trigger, Param: param, Context: ctx}
}

func newAuthRequest(authReq *odplugin.AuthRequest) *request {
	return &request{
		Kind: "provider",
		Name: authReq.ProviderName,
		Param: struct {
			Action   string                 `json:"action"`
			AuthData map[string]interface{} `json:"auth_data"`
		}{authReq.Action, authReq.AuthData},
	}
}

// TODO(limouren): reduce copying of this method
func (req *request) MarshalJSON() ([]byte, error) {
	pluginCtx := odplugin.ContextMap(req.Context)
	if rawParam, ok := req.Param.(json.RawMessage); ok {
		rawParamReq := struct {
			Kind    string                 `json:"kind"`
			Name    string                 `json:"name,omitempty"`
			Param   json.RawMessage        `json:"param,omitempty"`
			Context map[string]interface{} `json:"context,omitempty"`
		}{req.Kind, req.Name, rawParam, pluginCtx}
		return json.Marshal(&rawParamReq)
	}

	paramReq := struct {
		Kind    string                 `json:"kind"`
		Name    string                 `json:"name,omitempty"`
		Param   interface{}            `json:"param,omitempty"`
		Context map[string]interface{} `json:"context,omitempty"`
	}{req.Kind, req.Name, req.Param, pluginCtx}

	return json.Marshal(&paramReq)
}

func (p *zmqTransport) State() odplugin.TransportState {
	return p.state
}

func (p *zmqTransport) SetInitHandler(f odplugin.TransportInitHandler) {
	p.initHandler = f
}

func (p *zmqTransport) setState(state odplugin.TransportState) {
	if state != p.state {
		oldState := p.state
		p.state = state
		p.logger.Infof("Transport state changes from %v to %v.", oldState, p.state)
	}
}

func (p *zmqTransport) RequestInit() {
	for {
		address := <-p.broker.freshWorkers

		if p.state != odplugin.TransportStateUninitialized {
			// Although the plugin is only initialized once, we need
			// to clear the channel buffer so that broker doesn't get stuck
			continue
		}

		p.logger.Debugf("zmq transport got fresh worker %s", string(address))

		// TODO: Only send init to the new address. For now, we let
		// the broker decide.
		out, err := p.RunInit()
		if p.initHandler != nil {
			handlerError := p.initHandler(out, err)
			if err != nil || handlerError != nil {
				p.setState(odplugin.TransportStateError)
			}
		}
		p.setState(odplugin.TransportStateReady)
	}
}

func (p *zmqTransport) RunInit() (out []byte, err error) {
	req := request{Kind: "init", Timeout: initRequestTimeout}
	for {
		out, err = p.ipc(&req)
		if err == nil {
			break
		}
		p.logger.WithField("err", err).Warnf(`zmq/rpc: Unable to send init request to plugin "%s". Retrying...`, p.name)
	}
	return
}

func (p *zmqTransport) RunLambda(ctx context.Context, name string, in []byte) (out []byte, err error) {
	out, err = p.rpc(newLambdaRequest(ctx, name, in))
	return
}

func (p *zmqTransport) RunHandler(name string, in []byte) (out []byte, err error) {
	out, err = p.rpc(newHandlerRequest(name, in))
	return
}

func (p *zmqTransport) RunHook(ctx context.Context, recordType string, trigger string, record *skydb.Record, originalRecord *skydb.Record) (*skydb.Record, error) {
	out, err := p.rpc(newHookRequest(trigger, record, originalRecord, ctx))
	if err != nil {
		return nil, err
	}

	var recordout skydb.Record
	if err := json.Unmarshal(out, (*skyconv.JSONRecord)(&recordout)); err != nil {
		p.logger.WithField("data", string(out)).Error("failed to unmarshal record")
		return nil, fmt.Errorf("failed to unmarshal record: %v", err)
	}
	recordout.OwnerID = record.OwnerID
	recordout.CreatedAt = record.CreatedAt
	recordout.CreatorID = record.CreatorID
	recordout.UpdatedAt = record.UpdatedAt
	recordout.UpdaterID = record.UpdaterID

	return &recordout, nil
}

func (p *zmqTransport) RunTimer(name string, in []byte) (out []byte, err error) {
	req := request{Kind: "timer", Name: name}
	out, err = p.rpc(&req)
	return
}

func (p *zmqTransport) RunProvider(request *odplugin.AuthRequest) (resp *odplugin.AuthResponse, err error) {
	req := newAuthRequest(request)
	out, err := p.rpc(req)
	if err != nil {
		return
	}

	err = json.Unmarshal(out, &resp)
	return
}

func (p *zmqTransport) rpc(req *request) (out []byte, err error) {
	var rawResp []byte

	rawResp, err = p.ipc(req)
	if err != nil {
		return
	}

	var resp struct {
		Result json.RawMessage   `json:"result"`
		Err    *common.ExecError `json:"error"`
	}

	if err = json.Unmarshal(rawResp, &resp); err != nil {
		return
	}
	if resp.Err != nil {
		err = resp.Err
		return
	}

	out = resp.Result
	return
}

func (p *zmqTransport) ipc(req *request) (out []byte, err error) {
	var (
		in      []byte
		reqSock *goczmq.Sock
	)

	in, err = json.Marshal(req)
	if err != nil {
		return
	}

	reqSock, err = goczmq.NewReq(p.iaddr)
	if err != nil {
		return
	}
	if req.Timeout > 0 {
		reqSock.SetRcvtimeo(req.Timeout)
	}
	err = reqSock.SendMessage([][]byte{in})
	if err != nil {
		return
	}

	msg, err := reqSock.RecvMessage()
	if err != nil {
		return
	}

	if len(msg) != 1 {
		err = fmt.Errorf("malformed resp msg = %s", msg)
	} else {
		out = msg[0]
	}

	return
}

type zmqTransportFactory struct {
}

func (f zmqTransportFactory) Open(name string, args []string) (transport odplugin.Transport) {
	const internalAddrFmt = `inproc://%s`

	internalAddr := fmt.Sprintf(internalAddrFmt, name)
	externalAddr := args[0]

	broker, err := NewBroker(name, internalAddr, externalAddr)
	logger := log.WithFields(log.Fields{"plugin": name})
	if err != nil {
		logger.Panicf("Failed to init broker for zmq transport: %v", err)
	}

	p := zmqTransport{
		state:  odplugin.TransportStateUninitialized,
		name:   name,
		iaddr:  internalAddr,
		eaddr:  externalAddr,
		broker: broker,
		logger: logger,
	}

	go func() {
		logger.Infof("Running zmq broker:\niaddr = %s\neaddr = %s", internalAddr, externalAddr)
		broker.Run()
	}()

	return &p
}

func init() {
	// A new zmq socket is created here so that a new zmq context is created
	// for the process.
	// If the zmq context is not created first, it will be created at the
	// time when zmq sockets are created, which might cause zsys_init to
	// fail since it is not thread-safe.
	// goczmq does not provide function to init context, so we create
	// a new socket here and throw it away.
	router, err := goczmq.NewRouter(`inproc://init`)
	if err != nil {
		panic("unable to initialize zmq")
	}
	defer router.Destroy()

	// Register the zmq transport factory
	odplugin.RegisterTransport("zmq", zmqTransportFactory{})
}
<commit_msg>Hot fix to recover from zmq crash and log to errors<commit_after>package zmq

import (
	"encoding/json"
	"fmt"
	"runtime/debug"

	log "github.com/Sirupsen/logrus"
	odplugin "github.com/oursky/skygear/plugin"
	"github.com/oursky/skygear/plugin/common"
	"github.com/oursky/skygear/skyconv"
	"github.com/oursky/skygear/skydb"
	"github.com/zeromq/goczmq"
	"golang.org/x/net/context"
)

const initRequestTimeout = 2000

type zmqTransport struct {
	state       odplugin.TransportState
	name        string
	iaddr       string // the internal addr used by goroutines to make request to plugin
	eaddr       string // the addr exposed for plugin to connect to with REP.
	broker      *Broker
	initHandler odplugin.TransportInitHandler
	logger      *log.Entry
}

type request struct {
	Context context.Context
	Kind    string
	Name    string
	Param   interface{}
	Timeout int // timeout in millisecond
}

type hookRequest struct {
	Record   interface{} `json:"record"`
	Original interface{} `json:"original"`
}

// type-safe constructors for request.Param assignment

func newLambdaRequest(ctx context.Context, name string, args json.RawMessage) *request {
	return &request{Kind: "op", Name: name, Param: args, Context: ctx}
}

func newHandlerRequest(name string, input json.RawMessage) *request {
	return &request{Kind: "handler", Name: name, Param: input}
}

func newHookRequest(trigger string, record *skydb.Record, originalRecord *skydb.Record, ctx context.Context) *request {
	param := hookRequest{
		Record:   (*skyconv.JSONRecord)(record),
		Original: (*skyconv.JSONRecord)(originalRecord),
	}
	return &request{Kind: "hook", Name: trigger, Param: param, Context: ctx}
}

func newAuthRequest(authReq *odplugin.AuthRequest) *request {
	return &request{
		Kind: "provider",
		Name: authReq.ProviderName,
		Param: struct {
			Action   string                 `json:"action"`
			AuthData map[string]interface{} `json:"auth_data"`
		}{authReq.Action, authReq.AuthData},
	}
}

// TODO(limouren): reduce copying of this method
func (req *request) MarshalJSON() ([]byte, error) {
	pluginCtx := odplugin.ContextMap(req.Context)
	if rawParam, ok := req.Param.(json.RawMessage); ok {
		rawParamReq := struct {
			Kind    string                 `json:"kind"`
			Name    string                 `json:"name,omitempty"`
			Param   json.RawMessage        `json:"param,omitempty"`
			Context map[string]interface{} `json:"context,omitempty"`
		}{req.Kind, req.Name, rawParam, pluginCtx}
		return json.Marshal(&rawParamReq)
	}

	paramReq := struct {
		Kind    string                 `json:"kind"`
		Name    string                 `json:"name,omitempty"`
		Param   interface{}            `json:"param,omitempty"`
		Context map[string]interface{} `json:"context,omitempty"`
	}{req.Kind, req.Name, req.Param, pluginCtx}

	return json.Marshal(&paramReq)
}

func (p *zmqTransport) State() odplugin.TransportState {
	return p.state
}

func (p *zmqTransport) SetInitHandler(f odplugin.TransportInitHandler) {
	p.initHandler = f
}

func (p *zmqTransport) setState(state odplugin.TransportState) {
	if state != p.state {
		oldState := p.state
		p.state = state
		p.logger.Infof("Transport state changes from %v to %v.", oldState, p.state)
	}
}

func (p *zmqTransport) RequestInit() {
	for {
		address := <-p.broker.freshWorkers

		if p.state != odplugin.TransportStateUninitialized {
			// Although the plugin is only initialized once, we need
			// to clear the channel buffer so that broker doesn't get stuck
			continue
		}

		p.logger.Debugf("zmq transport got fresh worker %s", string(address))

		// TODO: Only send init to the new address. For now, we let
		// the broker decide.
		out, err := p.RunInit()
		if p.initHandler != nil {
			handlerError := p.initHandler(out, err)
			if err != nil || handlerError != nil {
				p.setState(odplugin.TransportStateError)
			}
		}
		p.setState(odplugin.TransportStateReady)
	}
}

func (p *zmqTransport) RunInit() (out []byte, err error) {
	req := request{Kind: "init", Timeout: initRequestTimeout}
	for {
		out, err = p.ipc(&req)
		if err == nil {
			break
		}
		p.logger.WithField("err", err).Warnf(`zmq/rpc: Unable to send init request to plugin "%s". Retrying...`, p.name)
	}
	return
}

func (p *zmqTransport) RunLambda(ctx context.Context, name string, in []byte) (out []byte, err error) {
	out, err = p.rpc(newLambdaRequest(ctx, name, in))
	return
}

func (p *zmqTransport) RunHandler(name string, in []byte) (out []byte, err error) {
	out, err = p.rpc(newHandlerRequest(name, in))
	return
}

func (p *zmqTransport) RunHook(ctx context.Context, recordType string, trigger string, record *skydb.Record, originalRecord *skydb.Record) (*skydb.Record, error) {
	out, err := p.rpc(newHookRequest(trigger, record, originalRecord, ctx))
	if err != nil {
		return nil, err
	}

	var recordout skydb.Record
	if err := json.Unmarshal(out, (*skyconv.JSONRecord)(&recordout)); err != nil {
		p.logger.WithField("data", string(out)).Error("failed to unmarshal record")
		return nil, fmt.Errorf("failed to unmarshal record: %v", err)
	}
	recordout.OwnerID = record.OwnerID
	recordout.CreatedAt = record.CreatedAt
	recordout.CreatorID = record.CreatorID
	recordout.UpdatedAt = record.UpdatedAt
	recordout.UpdaterID = record.UpdaterID

	return &recordout, nil
}

func (p *zmqTransport) RunTimer(name string, in []byte) (out []byte, err error) {
	req := request{Kind: "timer", Name: name}
	out, err = p.rpc(&req)
	return
}

func (p *zmqTransport) RunProvider(request *odplugin.AuthRequest) (resp *odplugin.AuthResponse, err error) {
	req := newAuthRequest(request)
	out, err := p.rpc(req)
	if err != nil {
		return
	}

	err = json.Unmarshal(out, &resp)
	return
}

func (p *zmqTransport) rpc(req *request) (out []byte, err error) {
	var rawResp []byte

	rawResp, err = p.ipc(req)
	if err != nil {
		return
	}

	var resp struct {
		Result json.RawMessage   `json:"result"`
		Err    *common.ExecError `json:"error"`
	}

	if err = json.Unmarshal(rawResp, &resp); err != nil {
		return
	}
	if resp.Err != nil {
		err = resp.Err
		return
	}

	out = resp.Result
	return
}

func (p *zmqTransport) ipc(req *request) (out []byte, err error) {
	defer func() {
		if r := recover(); r != nil {
			log.WithField("recovered", r).
				Errorln("panic occurred while calling plugin ipc")
			log.Errorf("%s", debug.Stack())
			return
		}
	}()
	var (
		in      []byte
		reqSock *goczmq.Sock
	)

	in, err = json.Marshal(req)
	if err != nil {
		return
	}

	reqSock, err = goczmq.NewReq(p.iaddr)
	if err != nil {
		return
	}
	if req.Timeout > 0 {
		reqSock.SetRcvtimeo(req.Timeout)
	}
	err = reqSock.SendMessage([][]byte{in})
	if err != nil {
		return
	}

	msg, err := reqSock.RecvMessage()
	if err != nil {
		return
	}

	if len(msg) != 1 {
		err = fmt.Errorf("malformed resp msg = %s", msg)
	} else {
		out = msg[0]
	}

	return
}

type zmqTransportFactory struct {
}

func (f zmqTransportFactory) Open(name string, args []string) (transport odplugin.Transport) {
	const internalAddrFmt = `inproc://%s`

	internalAddr := fmt.Sprintf(internalAddrFmt, name)
	externalAddr := args[0]

	broker, err := NewBroker(name, internalAddr, externalAddr)
	logger := log.WithFields(log.Fields{"plugin": name})
	if err != nil {
		logger.Panicf("Failed to init broker for zmq transport: %v", err)
	}

	p := zmqTransport{
		state:  odplugin.TransportStateUninitialized,
		name:   name,
		iaddr:  internalAddr,
		eaddr:  externalAddr,
		broker: broker,
		logger: logger,
	}

	go func() {
		logger.Infof("Running zmq broker:\niaddr = %s\neaddr = %s", internalAddr, externalAddr)
		broker.Run()
	}()

	return &p
}

func init() {
	// A new zmq socket is created here so that a new zmq context is created
	// for the process.
	// If the zmq context is not created first, it will be created at the
	// time when zmq sockets are created, which might cause zsys_init to
	// fail since it is not thread-safe.
	// goczmq does not provide function to init context, so we create
	// a new socket here and throw it away.
	router, err := goczmq.NewRouter(`inproc://init`)
	if err != nil {
		panic("unable to initialize zmq")
	}
	defer router.Destroy()

	// Register the zmq transport factory
	odplugin.RegisterTransport("zmq", zmqTransportFactory{})
}
<|endoftext|>


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

// This file records the static ranks of the locks in the runtime. If a lock
// is not given a rank, then it is assumed to be a leaf lock, which means no other
// lock can be acquired while it is held. Therefore, leaf locks do not need to be
// given an explicit rank. We list all of the architecture-independent leaf locks
// for documentation purposes, but don't list any of the architecture-dependent
// locks (which are all leaf locks). debugLock is ignored for ranking, since it is used
// when printing out lock ranking errors.
//
// lockInit(l *mutex, rank int) is used to set the rank of lock before it is used.
// If there is no clear place to initialize a lock, then the rank of a lock can be
// specified during the lock call itself via lockWithrank(l *mutex, rank int).
//
// Besides the static lock ranking (which is a total ordering of the locks), we
// also represent and enforce the actual partial order among the locks in the
// arcs[] array below. That is, if it is possible that lock B can be acquired when
// lock A is the previous acquired lock that is still held, then there should be
// an entry for A in arcs[B][]. We will currently fail not only if the total order
// (the lock ranking) is violated, but also if there is a missing entry in the
// partial order.

package runtime

type lockRank int

// Constants representing the lock rank of the architecture-independent locks in
// the runtime.
const (
	lockRankDummy lockRank = iota

	// Locks held above sched
	lockRankScavenge
	lockRankForcegc
	lockRankSweepWaiters
	lockRankAssistQueue
	lockRankCpuprof
	lockRankSweep

	lockRankSched
	lockRankDeadlock
	lockRankPanic
	lockRankAllg
	lockRankAllp
	lockRankPollDesc

	lockRankTimers // Multiple timers locked simultaneously in destroy()
	lockRankItab
	lockRankReflectOffs
	lockRankHchan // Multiple hchans acquired in lock order in syncadjustsudogs()
	lockRankFin
	lockRankNotifyList
	lockRankTraceBuf
	lockRankTraceStrings
	lockRankMspanSpecial
	lockRankProf
	lockRankGcBitsArenas
	lockRankRoot
	lockRankTrace
	lockRankTraceStackTab
	lockRankNetpollInit

	lockRankRwmutexW
	lockRankRwmutexR

	lockRankMcentral
	lockRankSpine
	lockRankStackpool
	lockRankStackLarge
	lockRankDefer
	lockRankSudog

	// Memory-related non-leaf locks
	lockRankWbufSpans
	lockRankMheap

	// Memory-related leaf locks
	lockRankMheapSpecial
	lockRankGlobalAlloc

	// Other leaf locks
	lockRankGFree

	// Leaf locks with no dependencies, so these constants are not actually used anywhere.
	// There are other architecture-dependent leaf locks as well.
	lockRankNewmHandoff
	lockRankDebugPtrmask
	lockRankFaketimeState
	lockRankTicks
	lockRankRaceFini
	lockRankPollCache
	lockRankDebug
)

// lockRankLeafRank is the rank of lock that does not have a declared rank, and hence is
// a leaf lock.
const lockRankLeafRank lockRank = 1000

// lockNames gives the names associated with each of the above ranks
var lockNames = []string{
	lockRankDummy: "",

	lockRankScavenge:     "scavenge",
	lockRankForcegc:      "forcegc",
	lockRankSweepWaiters: "sweepWaiters",
	lockRankAssistQueue:  "assistQueue",
	lockRankCpuprof:      "cpuprof",
	lockRankSweep:        "sweep",

	lockRankSched:    "sched",
	lockRankDeadlock: "deadlock",
	lockRankPanic:    "panic",
	lockRankAllg:     "allg",
	lockRankAllp:     "allp",
	lockRankPollDesc: "pollDesc",

	lockRankTimers:      "timers",
	lockRankItab:        "itab",
	lockRankReflectOffs: "reflectOffs",

	lockRankHchan:         "hchan",
	lockRankFin:           "fin",
	lockRankNotifyList:    "notifyList",
	lockRankTraceBuf:      "traceBuf",
	lockRankTraceStrings:  "traceStrings",
	lockRankMspanSpecial:  "mspanSpecial",
	lockRankProf:          "prof",
	lockRankGcBitsArenas:  "gcBitsArenas",
	lockRankRoot:          "root",
	lockRankTrace:         "trace",
	lockRankTraceStackTab: "traceStackTab",
	lockRankNetpollInit:   "netpollInit",

	lockRankRwmutexW: "rwmutexW",
	lockRankRwmutexR: "rwmutexR",

	lockRankMcentral:   "mcentral",
	lockRankSpine:      "spine",
	lockRankStackpool:  "stackpool",
	lockRankStackLarge: "stackLarge",
	lockRankDefer:      "defer",
	lockRankSudog:      "sudog",

	lockRankWbufSpans: "wbufSpans",
	lockRankMheap:     "mheap",

	lockRankMheapSpecial: "mheapSpecial",
	lockRankGlobalAlloc:  "globalAlloc.mutex",

	lockRankGFree: "gFree",

	lockRankNewmHandoff:   "newmHandoff.lock",
	lockRankDebugPtrmask:  "debugPtrmask.lock",
	lockRankFaketimeState: "faketimeState.lock",
	lockRankTicks:         "ticks.lock",
	lockRankRaceFini:      "raceFiniLock",
	lockRankPollCache:     "pollCache.lock",
	lockRankDebug:         "debugLock",
}

func (rank lockRank) String() string {
	if rank == 0 {
		return "UNKNOWN"
	}
	if rank == lockRankLeafRank {
		return "LEAF"
	}
	return lockNames[rank]
}

// lockPartialOrder is a partial order among the various lock types, listing the immediate
// ordering that has actually been observed in the runtime. Each entry (which
// corresponds to a particular lock rank) specifies the list of locks that can be
// already be held immediately "above" it.
//
// So, for example, the lockRankSched entry shows that all the locks preceding it in
// rank can actually be held. The fin lock shows that only the sched, timers, or
// hchan lock can be held immediately above it when it is acquired.
var lockPartialOrder [][]lockRank = [][]lockRank{
	lockRankDummy:         {},
	lockRankScavenge:      {},
	lockRankForcegc:       {},
	lockRankSweepWaiters:  {},
	lockRankAssistQueue:   {},
	lockRankCpuprof:       {},
	lockRankSweep:         {},
	lockRankSched:         {lockRankScavenge, lockRankForcegc, lockRankSweepWaiters, lockRankAssistQueue, lockRankCpuprof, lockRankSweep},
	lockRankDeadlock:      {lockRankDeadlock},
	lockRankPanic:         {lockRankDeadlock},
	lockRankAllg:          {lockRankSched, lockRankPanic},
	lockRankAllp:          {lockRankSched},
	lockRankPollDesc:      {},
	lockRankTimers:        {lockRankScavenge, lockRankSched, lockRankAllp, lockRankPollDesc, lockRankTimers},
	lockRankItab:          {},
	lockRankReflectOffs:   {lockRankItab},
	lockRankHchan:         {lockRankScavenge, lockRankSweep, lockRankHchan},
	lockRankFin:           {lockRankSched, lockRankAllg, lockRankTimers, lockRankHchan},
	lockRankNotifyList:    {},
	lockRankTraceBuf:      {},
	lockRankTraceStrings:  {lockRankTraceBuf},
	lockRankMspanSpecial:  {lockRankScavenge, lockRankCpuprof, lockRankSched, lockRankAllg, lockRankAllp, lockRankTimers, lockRankItab, lockRankReflectOffs, lockRankHchan, lockRankNotifyList, lockRankTraceBuf, lockRankTraceStrings},
	lockRankProf:          {lockRankScavenge, lockRankAssistQueue, lockRankCpuprof, lockRankSweep, lockRankSched, lockRankAllg, lockRankAllp, lockRankTimers, lockRankItab, lockRankReflectOffs, lockRankNotifyList, lockRankTraceBuf, lockRankTraceStrings, lockRankHchan},
	lockRankGcBitsArenas:  {lockRankScavenge, lockRankAssistQueue, lockRankCpuprof, lockRankSched, lockRankAllg, lockRankTimers, lockRankItab, lockRankReflectOffs, lockRankNotifyList, lockRankTraceBuf, lockRankTraceStrings, lockRankHchan},
	lockRankRoot:          {},
	lockRankTrace:         {lockRankScavenge, lockRankAssistQueue, lockRankSched, lockRankHchan, lockRankTraceBuf, lockRankTraceStrings, lockRankRoot, lockRankSweep},
	lockRankTraceStackTab: {lockRankScavenge, lockRankSweepWaiters, lockRankAssistQueue, lockRankSweep, lockRankSched, lockRankTimers, lockRankHchan, lockRankFin, lockRankNotifyList, lockRankTraceBuf, lockRankTraceStrings, lockRankRoot, lockRankTrace},
	lockRankNetpollInit:   {lockRankTimers},

	lockRankRwmutexW: {},
	lockRankRwmutexR: {lockRankRwmutexW},

	lockRankMcentral:     {lockRankScavenge, lockRankForcegc, lockRankAssistQueue, lockRankCpuprof, lockRankSweep, lockRankSched, lockRankAllg, lockRankAllp, lockRankTimers, lockRankItab, lockRankReflectOffs, lockRankNotifyList, lockRankTraceBuf, lockRankTraceStrings, lockRankHchan},
	lockRankSpine:        {lockRankScavenge, lockRankCpuprof, lockRankSched, lockRankAllg, lockRankTimers, lockRankItab, lockRankReflectOffs, lockRankNotifyList, lockRankTraceBuf, lockRankTraceStrings, lockRankHchan},
	lockRankStackpool:    {lockRankScavenge, lockRankSweepWaiters, lockRankAssistQueue, lockRankCpuprof, lockRankSweep, lockRankSched, lockRankPollDesc, lockRankTimers, lockRankItab, lockRankReflectOffs, lockRankHchan, lockRankFin, lockRankNotifyList, lockRankTraceBuf, lockRankTraceStrings, lockRankProf, lockRankGcBitsArenas, lockRankRoot, lockRankTrace, lockRankTraceStackTab, lockRankNetpollInit, lockRankRwmutexR, lockRankMcentral, lockRankSpine},
	lockRankStackLarge:   {lockRankAssistQueue, lockRankSched, lockRankItab, lockRankHchan, lockRankProf, lockRankGcBitsArenas, lockRankRoot, lockRankMcentral},
	lockRankDefer:        {},
	lockRankSudog:        {lockRankNotifyList, lockRankHchan},
	lockRankWbufSpans:    {lockRankScavenge, lockRankSweepWaiters, lockRankAssistQueue, lockRankSweep, lockRankSched, lockRankAllg, lockRankPollDesc, lockRankTimers, lockRankItab, lockRankReflectOffs, lockRankHchan, lockRankNotifyList, lockRankTraceStrings, lockRankMspanSpecial, lockRankProf, lockRankRoot, lockRankDefer, lockRankSudog},
	lockRankMheap:        {lockRankScavenge, lockRankSweepWaiters, lockRankAssistQueue, lockRankCpuprof, lockRankSweep, lockRankSched, lockRankAllg, lockRankAllp, lockRankPollDesc, lockRankTimers, lockRankItab, lockRankReflectOffs, lockRankNotifyList, lockRankTraceBuf, lockRankTraceStrings, lockRankHchan, lockRankMspanSpecial, lockRankProf, lockRankGcBitsArenas, lockRankRoot, lockRankMcentral, lockRankStackpool, lockRankStackLarge, lockRankDefer, lockRankSudog, lockRankWbufSpans},
	lockRankMheapSpecial: {lockRankScavenge, lockRankCpuprof, lockRankSched, lockRankAllg, lockRankAllp, lockRankTimers, lockRankItab, lockRankReflectOffs, lockRankNotifyList, lockRankTraceBuf, lockRankTraceStrings, lockRankHchan},
	lockRankGlobalAlloc:  {lockRankSpine, lockRankMheap},

	lockRankGFree: {lockRankSched},

	lockRankNewmHandoff:   {},
	lockRankDebugPtrmask:  {},
	lockRankFaketimeState: {},
	lockRankTicks:         {},
	lockRankRaceFini:      {},
	lockRankPollCache:     {},
	lockRankDebug:         {},
}
<commit_msg>runtime: explictly state lock ordering direction<commit_after>// Copyright 2020 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file records the static ranks of the locks in the runtime. If a lock
// is not given a rank, then it is assumed to be a leaf lock, which means no other
// lock can be acquired while it is held. Therefore, leaf locks do not need to be
// given an explicit rank. We list all of the architecture-independent leaf locks
// for documentation purposes, but don't list any of the architecture-dependent
// locks (which are all leaf locks). debugLock is ignored for ranking, since it is used
// when printing out lock ranking errors.
//
// lockInit(l *mutex, rank int) is used to set the rank of lock before it is used.
// If there is no clear place to initialize a lock, then the rank of a lock can be
// specified during the lock call itself via lockWithrank(l *mutex, rank int).
//
// Besides the static lock ranking (which is a total ordering of the locks), we
// also represent and enforce the actual partial order among the locks in the
// arcs[] array below. That is, if it is possible that lock B can be acquired when
// lock A is the previous acquired lock that is still held, then there should be
// an entry for A in arcs[B][]. We will currently fail not only if the total order
// (the lock ranking) is violated, but also if there is a missing entry in the
// partial order.

package runtime

type lockRank int

// Constants representing the lock rank of the architecture-independent locks in
// the runtime. Locks with lower rank must be taken before locks with higher
// rank.
const (
	lockRankDummy lockRank = iota

	// Locks held above sched
	lockRankScavenge
	lockRankForcegc
	lockRankSweepWaiters
	lockRankAssistQueue
	lockRankCpuprof
	lockRankSweep

	lockRankSched
	lockRankDeadlock
	lockRankPanic
	lockRankAllg
	lockRankAllp
	lockRankPollDesc

	lockRankTimers // Multiple timers locked simultaneously in destroy()
	lockRankItab
	lockRankReflectOffs
	lockRankHchan // Multiple hchans acquired in lock order in syncadjustsudogs()
	lockRankFin
	lockRankNotifyList
	lockRankTraceBuf
	lockRankTraceStrings
	lockRankMspanSpecial
	lockRankProf
	lockRankGcBitsArenas
	lockRankRoot
	lockRankTrace
	lockRankTraceStackTab
	lockRankNetpollInit

	lockRankRwmutexW
	lockRankRwmutexR

	lockRankMcentral
	lockRankSpine
	lockRankStackpool
	lockRankStackLarge
	lockRankDefer
	lockRankSudog

	// Memory-related non-leaf locks
	lockRankWbufSpans
	lockRankMheap

	// Memory-related leaf locks
	lockRankMheapSpecial
	lockRankGlobalAlloc

	// Other leaf locks
	lockRankGFree

	// Leaf locks with no dependencies, so these constants are not actually used anywhere.
	// There are other architecture-dependent leaf locks as well.
	lockRankNewmHandoff
	lockRankDebugPtrmask
	lockRankFaketimeState
	lockRankTicks
	lockRankRaceFini
	lockRankPollCache
	lockRankDebug
)

// lockRankLeafRank is the rank of lock that does not have a declared rank, and hence is
// a leaf lock.
const lockRankLeafRank lockRank = 1000

// lockNames gives the names associated with each of the above ranks
var lockNames = []string{
	lockRankDummy: "",

	lockRankScavenge:     "scavenge",
	lockRankForcegc:      "forcegc",
	lockRankSweepWaiters: "sweepWaiters",
	lockRankAssistQueue:  "assistQueue",
	lockRankCpuprof:      "cpuprof",
	lockRankSweep:        "sweep",

	lockRankSched:    "sched",
	lockRankDeadlock: "deadlock",
	lockRankPanic:    "panic",
	lockRankAllg:     "allg",
	lockRankAllp:     "allp",
	lockRankPollDesc: "pollDesc",

	lockRankTimers:      "timers",
	lockRankItab:        "itab",
	lockRankReflectOffs: "reflectOffs",

	lockRankHchan:         "hchan",
	lockRankFin:           "fin",
	lockRankNotifyList:    "notifyList",
	lockRankTraceBuf:      "traceBuf",
	lockRankTraceStrings:  "traceStrings",
	lockRankMspanSpecial:  "mspanSpecial",
	lockRankProf:          "prof",
	lockRankGcBitsArenas:  "gcBitsArenas",
	lockRankRoot:          "root",
	lockRankTrace:         "trace",
	lockRankTraceStackTab: "traceStackTab",
	lockRankNetpollInit:   "netpollInit",

	lockRankRwmutexW: "rwmutexW",
	lockRankRwmutexR: "rwmutexR",

	lockRankMcentral:   "mcentral",
	lockRankSpine:      "spine",
	lockRankStackpool:  "stackpool",
	lockRankStackLarge: "stackLarge",
	lockRankDefer:      "defer",
	lockRankSudog:      "sudog",

	lockRankWbufSpans: "wbufSpans",
	lockRankMheap:     "mheap",

	lockRankMheapSpecial: "mheapSpecial",
	lockRankGlobalAlloc:  "globalAlloc.mutex",

	lockRankGFree: "gFree",

	lockRankNewmHandoff:   "newmHandoff.lock",
	lockRankDebugPtrmask:  "debugPtrmask.lock",
	lockRankFaketimeState: "faketimeState.lock",
	lockRankTicks:         "ticks.lock",
	lockRankRaceFini:      "raceFiniLock",
	lockRankPollCache:     "pollCache.lock",
	lockRankDebug:         "debugLock",
}

func (rank lockRank) String() string {
	if rank == 0 {
		return "UNKNOWN"
	}
	if rank == lockRankLeafRank {
		return "LEAF"
	}
	return lockNames[rank]
}

// lockPartialOrder is a partial order among the various lock types, listing the immediate
// ordering that has actually been observed in the runtime. Each entry (which
// corresponds to a particular lock rank) specifies the list of locks that can be
// already be held immediately "above" it.
//
// So, for example, the lockRankSched entry shows that all the locks preceding it in
// rank can actually be held. The fin lock shows that only the sched, timers, or
// hchan lock can be held immediately above it when it is acquired.
var lockPartialOrder [][]lockRank = [][]lockRank{
	lockRankDummy:         {},
	lockRankScavenge:      {},
	lockRankForcegc:       {},
	lockRankSweepWaiters:  {},
	lockRankAssistQueue:   {},
	lockRankCpuprof:       {},
	lockRankSweep:         {},
	lockRankSched:         {lockRankScavenge, lockRankForcegc, lockRankSweepWaiters, lockRankAssistQueue, lockRankCpuprof, lockRankSweep},
	lockRankDeadlock:      {lockRankDeadlock},
	lockRankPanic:         {lockRankDeadlock},
	lockRankAllg:          {lockRankSched, lockRankPanic},
	lockRankAllp:          {lockRankSched},
	lockRankPollDesc:      {},
	lockRankTimers:        {lockRankScavenge, lockRankSched, lockRankAllp, lockRankPollDesc, lockRankTimers},
	lockRankItab:          {},
	lockRankReflectOffs:   {lockRankItab},
	lockRankHchan:         {lockRankScavenge, lockRankSweep, lockRankHchan},
	lockRankFin:           {lockRankSched, lockRankAllg, lockRankTimers, lockRankHchan},
	lockRankNotifyList:    {},
	lockRankTraceBuf:      {},
	lockRankTraceStrings:  {lockRankTraceBuf},
	lockRankMspanSpecial:  {lockRankScavenge, lockRankCpuprof, lockRankSched, lockRankAllg, lockRankAllp, lockRankTimers, lockRankItab, lockRankReflectOffs, lockRankHchan, lockRankNotifyList, lockRankTraceBuf, lockRankTraceStrings},
	lockRankProf:          {lockRankScavenge, lockRankAssistQueue, lockRankCpuprof, lockRankSweep, lockRankSched, lockRankAllg, lockRankAllp, lockRankTimers, lockRankItab, lockRankReflectOffs, lockRankNotifyList, lockRankTraceBuf, lockRankTraceStrings, lockRankHchan},
	lockRankGcBitsArenas:  {lockRankScavenge, lockRankAssistQueue, lockRankCpuprof, lockRankSched, lockRankAllg, lockRankTimers, lockRankItab, lockRankReflectOffs, lockRankNotifyList, lockRankTraceBuf, lockRankTraceStrings, lockRankHchan},
	lockRankRoot:          {},
	lockRankTrace:         {lockRankScavenge, lockRankAssistQueue, lockRankSched, lockRankHchan, lockRankTraceBuf, lockRankTraceStrings, lockRankRoot, lockRankSweep},
	lockRankTraceStackTab: {lockRankScavenge, lockRankSweepWaiters, lockRankAssistQueue, lockRankSweep, lockRankSched, lockRankTimers, lockRankHchan, lockRankFin, lockRankNotifyList, lockRankTraceBuf, lockRankTraceStrings, lockRankRoot, lockRankTrace},
	lockRankNetpollInit:   {lockRankTimers},

	lockRankRwmutexW: {},
	lockRankRwmutexR: {lockRankRwmutexW},

	lockRankMcentral:     {lockRankScavenge, lockRankForcegc, lockRankAssistQueue, lockRankCpuprof, lockRankSweep, lockRankSched, lockRankAllg, lockRankAllp, lockRankTimers, lockRankItab, lockRankReflectOffs, lockRankNotifyList, lockRankTraceBuf, lockRankTraceStrings, lockRankHchan},
	lockRankSpine:        {lockRankScavenge, lockRankCpuprof, lockRankSched, lockRankAllg, lockRankTimers, lockRankItab, lockRankReflectOffs, lockRankNotifyList, lockRankTraceBuf, lockRankTraceStrings, lockRankHchan},
	lockRankStackpool:    {lockRankScavenge, lockRankSweepWaiters, lockRankAssistQueue, lockRankCpuprof, lockRankSweep, lockRankSched, lockRankPollDesc, lockRankTimers, lockRankItab, lockRankReflectOffs, lockRankHchan, lockRankFin, lockRankNotifyList, lockRankTraceBuf, lockRankTraceStrings, lockRankProf, lockRankGcBitsArenas, lockRankRoot, lockRankTrace, lockRankTraceStackTab, lockRankNetpollInit, lockRankRwmutexR, lockRankMcentral, lockRankSpine},
	lockRankStackLarge:   {lockRankAssistQueue, lockRankSched, lockRankItab, lockRankHchan, lockRankProf, lockRankGcBitsArenas, lockRankRoot, lockRankMcentral},
	lockRankDefer:        {},
	lockRankSudog:        {lockRankNotifyList, lockRankHchan},
	lockRankWbufSpans:    {lockRankScavenge, lockRankSweepWaiters, lockRankAssistQueue, lockRankSweep, lockRankSched, lockRankAllg, lockRankPollDesc, lockRankTimers, lockRankItab, lockRankReflectOffs, lockRankHchan, lockRankNotifyList, lockRankTraceStrings, lockRankMspanSpecial, lockRankProf, lockRankRoot, lockRankDefer, lockRankSudog},
	lockRankMheap:        {lockRankScavenge, lockRankSweepWaiters, lockRankAssistQueue, lockRankCpuprof, lockRankSweep, lockRankSched, lockRankAllg, lockRankAllp, lockRankPollDesc, lockRankTimers, lockRankItab, lockRankReflectOffs, lockRankNotifyList, lockRankTraceBuf, lockRankTraceStrings, lockRankHchan, lockRankMspanSpecial, lockRankProf, lockRankGcBitsArenas, lockRankRoot, lockRankMcentral, lockRankStackpool, lockRankStackLarge, lockRankDefer, lockRankSudog, lockRankWbufSpans},
	lockRankMheapSpecial: {lockRankScavenge, lockRankCpuprof, lockRankSched, lockRankAllg, lockRankAllp, lockRankTimers, lockRankItab, lockRankReflectOffs, lockRankNotifyList, lockRankTraceBuf, lockRankTraceStrings, lockRankHchan},
	lockRankGlobalAlloc:  {lockRankSpine, lockRankMheap},

	lockRankGFree: {lockRankSched},

	lockRankNewmHandoff:   {},
	lockRankDebugPtrmask:  {},
	lockRankFaketimeState: {},
	lockRankTicks:         {},
	lockRankRaceFini:      {},
	lockRankPollCache:     {},
	lockRankDebug:         {},
}
<|endoftext|>


<commit_before>package shells

import (
	"path"
	"path/filepath"
	"strconv"
	"strings"

	"errors"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/common"
)

type AbstractShell struct {
}

func (b *AbstractShell) GetFeatures(features *common.FeaturesInfo) {
	features.Artifacts = true
	features.Cache = true
}

func (b *AbstractShell) GetSupportedOptions() []string {
	return []string{"artifacts", "cache", "dependencies", "after_script"}
}

func (b *AbstractShell) writeCdBuildDir(w ShellWriter, info common.ShellScriptInfo) {
	w.Cd(info.Build.FullProjectDir())
}

func (b *AbstractShell) writeExports(w ShellWriter, info common.ShellScriptInfo) {
	for _, variable := range info.Build.GetAllVariables() {
		w.Variable(variable)
	}
}

func (b *AbstractShell) writeTLSCAInfo(w ShellWriter, build *common.Build, key string) {
	if build.TLSCAChain != "" {
		w.Variable(common.BuildVariable{
			Key:      key,
			Value:    build.TLSCAChain,
			Public:   true,
			Internal: true,
			File:     true,
		})
	}
}

func (b *AbstractShell) writeCloneCmd(w ShellWriter, build *common.Build, projectDir string) {
	templateDir := w.MkTmpDir("git-template")
	args := []string{"clone", build.RepoURL, projectDir, "--template", templateDir}

	w.RmDir(projectDir)
	w.Command("git", "config", "-f", path.Join(templateDir, "config"), "fetch.recurseSubmodules", "false")

	if depth := build.GetGitDepth(); depth != "" {
		w.Notice("Cloning repository for %s with git depth set to %s...", build.RefName, depth)
		args = append(args, "--depth", depth, "--branch", build.RefName)
	} else {
		w.Notice("Cloning repository...")
	}

	w.Command("git", args...)
	w.Cd(projectDir)
}

func (b *AbstractShell) writeFetchCmd(w ShellWriter, build *common.Build, projectDir string, gitDir string) {
	depth := build.GetGitDepth()

	w.IfDirectory(gitDir)
	if depth != "" {
		w.Notice("Fetching changes for %s with git depth set to %s...", build.RefName, depth)
	} else {
		w.Notice("Fetching changes...")
	}
	w.Cd(projectDir)
	w.Command("git", "config", "fetch.recurseSubmodules", "false")
	w.Command("git", "clean", "-ffdx")
	w.Command("git", "reset", "--hard")
	w.Command("git", "remote", "set-url", "origin", build.RepoURL)
	if depth != "" {
		var refspec string
		if build.Tag {
			refspec = "+refs/tags/" + build.RefName + ":refs/tags/" + build.RefName
		} else {
			refspec = "+refs/heads/" + build.RefName + ":refs/remotes/origin/" + build.RefName
		}
		w.Command("git", "fetch", "--depth", depth, "origin", "--prune", refspec)
	} else {
		w.Command("git", "fetch", "origin", "--prune", "+refs/heads/*:refs/remotes/origin/*", "+refs/tags/*:refs/tags/*")
	}
	w.Else()
	b.writeCloneCmd(w, build, projectDir)
	w.EndIf()
}

func (b *AbstractShell) writeCheckoutCmd(w ShellWriter, build *common.Build) {
	w.Notice("Checking out %s as %s...", build.Sha[0:8], build.RefName)
	// We remove a git index file, this is required if `git checkout` is terminated
	w.RmFile(".git/index.lock")
	w.Command("git", "checkout", "-q", build.Sha)
}

func (b *AbstractShell) cacheFile(build *common.Build, userKey string) (key, file string) {
	if build.CacheDir == "" {
		return
	}

	// Deduce cache key
	key = path.Join(build.Name, build.RefName)
	if userKey != "" {
		key = build.GetAllVariables().ExpandValue(userKey)
	}

	// Ignore cache without the key
	if key == "" {
		return
	}

	file = path.Join(build.CacheDir, key, "cache.zip")
	file, err := filepath.Rel(build.BuildDir, file)
	if err != nil {
		return "", ""
	}
	return
}

func (o *archivingOptions) CommandArguments() (args []string) {
	for _, path := range o.Paths {
		args = append(args, "--path", path)
	}

	if o.Untracked {
		args = append(args, "--untracked")
	}
	return
}

func (b *AbstractShell) guardRunnerCommand(w ShellWriter, runnerCommand string, action string, f func()) {
	if runnerCommand == "" {
		w.Warning("%s is not supported by this executor.", action)
		return
	}

	w.IfCmd(runnerCommand, "--version")
	f()
	w.Else()
	w.Warning("Missing %s. %s is disabled.", runnerCommand, action)
	w.EndIf()
}

func (b *AbstractShell) cacheExtractor(w ShellWriter, options *archivingOptions, info common.ShellScriptInfo) {
	if options == nil {
		return
	}

	// Skip restoring cache if no cache is defined
	if archiverArgs := options.CommandArguments(); len(archiverArgs) == 0 {
		return
	}

	// Skip archiving if no cache is defined
	cacheKey, cacheFile := b.cacheFile(info.Build, options.Key)
	if cacheKey == "" {
		return
	}

	args := []string{
		"cache-extractor",
		"--file", cacheFile,
	}

	// Generate cache download address
	if url := getCacheDownloadURL(info.Build, cacheKey); url != nil {
		args = append(args, "--url", url.String())
	}

	// Execute archive command
	b.guardRunnerCommand(w, info.RunnerCommand, "Extracting cache", func() {
		w.Notice("Checking cache for %s...", cacheKey)
		w.Command(info.RunnerCommand, args...)
	})
}

func (b *AbstractShell) downloadArtifacts(w ShellWriter, build *common.BuildInfo, info common.ShellScriptInfo) {
	args := []string{
		"artifacts-downloader",
		"--url",
		info.Build.Runner.URL,
		"--token",
		build.Token,
		"--id",
		strconv.Itoa(build.ID),
	}

	w.Notice("Downloading artifacts for %s (%d)...", build.Name, build.ID)
	w.Command(info.RunnerCommand, args...)
}

func (b *AbstractShell) buildArtifacts(dependencies *dependencies, info common.ShellScriptInfo) (otherBuilds []common.BuildInfo) {
	for _, otherBuild := range info.Build.DependsOnBuilds {
		if otherBuild.Artifacts == nil || otherBuild.Artifacts.Filename == "" {
			continue
		}
		if !dependencies.IsDependent(otherBuild.Name) {
			continue
		}
		otherBuilds = append(otherBuilds, otherBuild)
	}
	return
}

func (b *AbstractShell) downloadAllArtifacts(w ShellWriter, dependencies *dependencies, info common.ShellScriptInfo) {
	otherBuilds := b.buildArtifacts(dependencies, info)
	if len(otherBuilds) == 0 {
		return
	}

	b.guardRunnerCommand(w, info.RunnerCommand, "Artifacts downloading", func() {
		for _, otherBuild := range otherBuilds {
			b.downloadArtifacts(w, &otherBuild, info)
		}
	})
}

func (b *AbstractShell) writePrepareScript(w ShellWriter, info common.ShellScriptInfo) (err error) {
	b.writeExports(w, info)

	build := info.Build
	projectDir := build.FullProjectDir()
	gitDir := path.Join(build.FullProjectDir(), ".git")

	b.writeTLSCAInfo(w, info.Build, "GIT_SSL_CAINFO")
	b.writeTLSCAInfo(w, info.Build, "CI_SERVER_TLS_CA_FILE")

	if info.PreCloneScript != "" {
		b.writeCommands(w, info.PreCloneScript)
	}

	switch info.Build.GetGitStrategy() {
	case common.GitFetch:
		b.writeFetchCmd(w, build, projectDir, gitDir)

	case common.GitClone:
		b.writeCloneCmd(w, build, projectDir)

	default:
		return errors.New("unknown GIT_STRATEGY")
	}

	b.writeCheckoutCmd(w, build)

	// Parse options
	var options shellOptions
	err = info.Build.Options.Decode(&options)
	if err != nil {
		return
	}

	// Try to restore from main cache, if not found cache for master
	b.cacheExtractor(w, options.Cache, info)

	// Process all artifacts
	b.downloadAllArtifacts(w, options.Dependencies, info)
	return nil
}

// Write the given string of commands using the provided ShellWriter object.
func (b *AbstractShell) writeCommands(w ShellWriter, commands string) {
	commands = strings.TrimSpace(commands)
	for _, command := range strings.Split(commands, "\n") {
		command = strings.TrimSpace(command)
		if command != "" {
			w.Notice("$ %s", command)
		} else {
			w.EmptyLine()
		}
		w.Line(command)
		w.CheckForErrors()
	}
}

func (b *AbstractShell) writeBuildScript(w ShellWriter, info common.ShellScriptInfo) (err error) {
	b.writeExports(w, info)
	b.writeCdBuildDir(w, info)

	if info.PreBuildScript != "" {
		b.writeCommands(w, info.PreBuildScript)
	}

	commands := info.Build.Commands
	b.writeCommands(w, commands)

	return nil
}

func (b *AbstractShell) cacheArchiver(w ShellWriter, options *archivingOptions, info common.ShellScriptInfo) {
	if options == nil {
		return
	}

	// Skip archiving if no cache is defined
	cacheKey, cacheFile := b.cacheFile(info.Build, options.Key)
	if cacheKey == "" {
		return
	}

	args := []string{
		"cache-archiver",
		"--file", cacheFile,
	}

	// Create list of files to archive
	archiverArgs := options.CommandArguments()
	if len(archiverArgs) == 0 {
		// Skip creating archive
		return
	}
	args = append(args, archiverArgs...)

	// Generate cache upload address
	if url := getCacheUploadURL(info.Build, cacheKey); url != nil {
		args = append(args, "--url", url.String())
	}

	b.guardRunnerCommand(w, info.RunnerCommand, "Creating cache", func() {
		// Execute archive command
		w.Notice("Creating cache %s...", cacheKey)
		w.Command(info.RunnerCommand, args...)
	})
}

func (b *AbstractShell) uploadArtifacts(w ShellWriter, options *archivingOptions, info common.ShellScriptInfo) {
	if options == nil {
		return
	}
	if info.Build.Runner.URL == "" {
		return
	}

	args := []string{
		"artifacts-uploader",
		"--url",
		info.Build.Runner.URL,
		"--token",
		info.Build.Token,
		"--id",
		strconv.Itoa(info.Build.ID),
	}

	// Create list of files to archive
	archiverArgs := options.CommandArguments()
	if len(archiverArgs) == 0 {
		// Skip creating archive
		return
	}
	args = append(args, archiverArgs...)

	// Get artifacts:name
	if name, ok := info.Build.Options.GetString("artifacts", "name"); ok && name != "" {
		args = append(args, "--name", name)
	}

	// Get artifacts:expire_in
	if expireIn, ok := info.Build.Options.GetString("artifacts", "expire_in"); ok && expireIn != "" {
		args = append(args, "--expire-in", expireIn)
	}

	b.guardRunnerCommand(w, info.RunnerCommand, "Uploading artifacts", func() {
		w.Notice("Uploading artifacts...")
		w.Command(info.RunnerCommand, args...)
	})
}

func (b *AbstractShell) writeAfterScript(w ShellWriter, info common.ShellScriptInfo) error {
	shellOptions := struct {
		AfterScript []string `json:"after_script"`
	}{}
	err := info.Build.Options.Decode(&shellOptions)
	if err != nil {
		return err
	}

	if len(shellOptions.AfterScript) == 0 {
		return nil
	}

	b.writeExports(w, info)
	b.writeCdBuildDir(w, info)

	w.Notice("Running after script...")

	for _, command := range shellOptions.AfterScript {
		command = strings.TrimSpace(command)
		if command != "" {
			w.Notice("$ %s", command)
		} else {
			w.EmptyLine()
		}
		w.Line(command)
		w.CheckForErrors()
	}

	return nil
}

func (b *AbstractShell) writeArchiveCacheScript(w ShellWriter, info common.ShellScriptInfo) (err error) {
	// Parse options
	var options shellOptions
	err = info.Build.Options.Decode(&options)
	if err != nil {
		return
	}

	b.writeExports(w, info)
	b.writeCdBuildDir(w, info)
	b.writeTLSCAInfo(w, info.Build, "CI_SERVER_TLS_CA_FILE")

	// Find cached files and archive them
	b.cacheArchiver(w, options.Cache, info)
	return
}

func (b *AbstractShell) writeUploadArtifactsScript(w ShellWriter, info common.ShellScriptInfo) (err error) {
	// Parse options
	var options shellOptions
	err = info.Build.Options.Decode(&options)
	if err != nil {
		return
	}

	b.writeExports(w, info)
	b.writeCdBuildDir(w, info)
	b.writeTLSCAInfo(w, info.Build, "CI_SERVER_TLS_CA_FILE")

	// Upload artifacts
	b.uploadArtifacts(w, options.Artifacts, info)
	return
}

func (b *AbstractShell) writeScript(w ShellWriter, scriptType common.ShellScriptType, info common.ShellScriptInfo) (err error) {
	switch scriptType {
	case common.ShellPrepareScript:
		return b.writePrepareScript(w, info)

	case common.ShellBuildScript:
		return b.writeBuildScript(w, info)

	case common.ShellAfterScript:
		return b.writeAfterScript(w, info)

	case common.ShellArchiveCache:
		return b.writeArchiveCacheScript(w, info)

	case common.ShellUploadArtifacts:
		return b.writeUploadArtifactsScript(w, info)

	default:
		return errors.New("Not supported script type: " + string(scriptType))
	}
}
<commit_msg>Remove ./git/index.lock before fetching<commit_after>package shells

import (
	"path"
	"path/filepath"
	"strconv"
	"strings"

	"errors"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/common"
)

type AbstractShell struct {
}

func (b *AbstractShell) GetFeatures(features *common.FeaturesInfo) {
	features.Artifacts = true
	features.Cache = true
}

func (b *AbstractShell) GetSupportedOptions() []string {
	return []string{"artifacts", "cache", "dependencies", "after_script"}
}

func (b *AbstractShell) writeCdBuildDir(w ShellWriter, info common.ShellScriptInfo) {
	w.Cd(info.Build.FullProjectDir())
}

func (b *AbstractShell) writeExports(w ShellWriter, info common.ShellScriptInfo) {
	for _, variable := range info.Build.GetAllVariables() {
		w.Variable(variable)
	}
}

func (b *AbstractShell) writeTLSCAInfo(w ShellWriter, build *common.Build, key string) {
	if build.TLSCAChain != "" {
		w.Variable(common.BuildVariable{
			Key:      key,
			Value:    build.TLSCAChain,
			Public:   true,
			Internal: true,
			File:     true,
		})
	}
}

func (b *AbstractShell) writeCloneCmd(w ShellWriter, build *common.Build, projectDir string) {
	templateDir := w.MkTmpDir("git-template")
	args := []string{"clone", build.RepoURL, projectDir, "--template", templateDir}

	w.RmDir(projectDir)
	w.Command("git", "config", "-f", path.Join(templateDir, "config"), "fetch.recurseSubmodules", "false")

	if depth := build.GetGitDepth(); depth != "" {
		w.Notice("Cloning repository for %s with git depth set to %s...", build.RefName, depth)
		args = append(args, "--depth", depth, "--branch", build.RefName)
	} else {
		w.Notice("Cloning repository...")
	}

	w.Command("git", args...)
	w.Cd(projectDir)
}

func (b *AbstractShell) writeFetchCmd(w ShellWriter, build *common.Build, projectDir string, gitDir string) {
	depth := build.GetGitDepth()

	w.IfDirectory(gitDir)
	if depth != "" {
		w.Notice("Fetching changes for %s with git depth set to %s...", build.RefName, depth)
	} else {
		w.Notice("Fetching changes...")
	}
	w.Cd(projectDir)
	w.Command("git", "config", "fetch.recurseSubmodules", "false")

	// Remove existing .git/index.lock file which can fail the fetch command
	// The file can be left if previous build was terminated during git operation
	w.RmFile(".git/index.lock")

	w.Command("git", "clean", "-ffdx")
	w.Command("git", "reset", "--hard")
	w.Command("git", "remote", "set-url", "origin", build.RepoURL)
	if depth != "" {
		var refspec string
		if build.Tag {
			refspec = "+refs/tags/" + build.RefName + ":refs/tags/" + build.RefName
		} else {
			refspec = "+refs/heads/" + build.RefName + ":refs/remotes/origin/" + build.RefName
		}
		w.Command("git", "fetch", "--depth", depth, "origin", "--prune", refspec)
	} else {
		w.Command("git", "fetch", "origin", "--prune", "+refs/heads/*:refs/remotes/origin/*", "+refs/tags/*:refs/tags/*")
	}
	w.Else()
	b.writeCloneCmd(w, build, projectDir)
	w.EndIf()
}

func (b *AbstractShell) writeCheckoutCmd(w ShellWriter, build *common.Build) {
	w.Notice("Checking out %s as %s...", build.Sha[0:8], build.RefName)
	w.Command("git", "checkout", "-q", build.Sha)
}

func (b *AbstractShell) cacheFile(build *common.Build, userKey string) (key, file string) {
	if build.CacheDir == "" {
		return
	}

	// Deduce cache key
	key = path.Join(build.Name, build.RefName)
	if userKey != "" {
		key = build.GetAllVariables().ExpandValue(userKey)
	}

	// Ignore cache without the key
	if key == "" {
		return
	}

	file = path.Join(build.CacheDir, key, "cache.zip")
	file, err := filepath.Rel(build.BuildDir, file)
	if err != nil {
		return "", ""
	}
	return
}

func (o *archivingOptions) CommandArguments() (args []string) {
	for _, path := range o.Paths {
		args = append(args, "--path", path)
	}

	if o.Untracked {
		args = append(args, "--untracked")
	}
	return
}

func (b *AbstractShell) guardRunnerCommand(w ShellWriter, runnerCommand string, action string, f func()) {
	if runnerCommand == "" {
		w.Warning("%s is not supported by this executor.", action)
		return
	}

	w.IfCmd(runnerCommand, "--version")
	f()
	w.Else()
	w.Warning("Missing %s. %s is disabled.", runnerCommand, action)
	w.EndIf()
}

func (b *AbstractShell) cacheExtractor(w ShellWriter, options *archivingOptions, info common.ShellScriptInfo) {
	if options == nil {
		return
	}

	// Skip restoring cache if no cache is defined
	if archiverArgs := options.CommandArguments(); len(archiverArgs) == 0 {
		return
	}

	// Skip archiving if no cache is defined
	cacheKey, cacheFile := b.cacheFile(info.Build, options.Key)
	if cacheKey == "" {
		return
	}

	args := []string{
		"cache-extractor",
		"--file", cacheFile,
	}

	// Generate cache download address
	if url := getCacheDownloadURL(info.Build, cacheKey); url != nil {
		args = append(args, "--url", url.String())
	}

	// Execute archive command
	b.guardRunnerCommand(w, info.RunnerCommand, "Extracting cache", func() {
		w.Notice("Checking cache for %s...", cacheKey)
		w.Command(info.RunnerCommand, args...)
	})
}

func (b *AbstractShell) downloadArtifacts(w ShellWriter, build *common.BuildInfo, info common.ShellScriptInfo) {
	args := []string{
		"artifacts-downloader",
		"--url",
		info.Build.Runner.URL,
		"--token",
		build.Token,
		"--id",
		strconv.Itoa(build.ID),
	}

	w.Notice("Downloading artifacts for %s (%d)...", build.Name, build.ID)
	w.Command(info.RunnerCommand, args...)
}

func (b *AbstractShell) buildArtifacts(dependencies *dependencies, info common.ShellScriptInfo) (otherBuilds []common.BuildInfo) {
	for _, otherBuild := range info.Build.DependsOnBuilds {
		if otherBuild.Artifacts == nil || otherBuild.Artifacts.Filename == "" {
			continue
		}
		if !dependencies.IsDependent(otherBuild.Name) {
			continue
		}
		otherBuilds = append(otherBuilds, otherBuild)
	}
	return
}

func (b *AbstractShell) downloadAllArtifacts(w ShellWriter, dependencies *dependencies, info common.ShellScriptInfo) {
	otherBuilds := b.buildArtifacts(dependencies, info)
	if len(otherBuilds) == 0 {
		return
	}

	b.guardRunnerCommand(w, info.RunnerCommand, "Artifacts downloading", func() {
		for _, otherBuild := range otherBuilds {
			b.downloadArtifacts(w, &otherBuild, info)
		}
	})
}

func (b *AbstractShell) writePrepareScript(w ShellWriter, info common.ShellScriptInfo) (err error) {
	b.writeExports(w, info)

	build := info.Build
	projectDir := build.FullProjectDir()
	gitDir := path.Join(build.FullProjectDir(), ".git")

	b.writeTLSCAInfo(w, info.Build, "GIT_SSL_CAINFO")
	b.writeTLSCAInfo(w, info.Build, "CI_SERVER_TLS_CA_FILE")

	if info.PreCloneScript != "" {
		b.writeCommands(w, info.PreCloneScript)
	}

	switch info.Build.GetGitStrategy() {
	case common.GitFetch:
		b.writeFetchCmd(w, build, projectDir, gitDir)

	case common.GitClone:
		b.writeCloneCmd(w, build, projectDir)

	default:
		return errors.New("unknown GIT_STRATEGY")
	}

	b.writeCheckoutCmd(w, build)

	// Parse options
	var options shellOptions
	err = info.Build.Options.Decode(&options)
	if err != nil {
		return
	}

	// Try to restore from main cache, if not found cache for master
	b.cacheExtractor(w, options.Cache, info)

	// Process all artifacts
	b.downloadAllArtifacts(w, options.Dependencies, info)
	return nil
}

// Write the given string of commands using the provided ShellWriter object.
func (b *AbstractShell) writeCommands(w ShellWriter, commands string) {
	commands = strings.TrimSpace(commands)
	for _, command := range strings.Split(commands, "\n") {
		command = strings.TrimSpace(command)
		if command != "" {
			w.Notice("$ %s", command)
		} else {
			w.EmptyLine()
		}
		w.Line(command)
		w.CheckForErrors()
	}
}

func (b *AbstractShell) writeBuildScript(w ShellWriter, info common.ShellScriptInfo) (err error) {
	b.writeExports(w, info)
	b.writeCdBuildDir(w, info)

	if info.PreBuildScript != "" {
		b.writeCommands(w, info.PreBuildScript)
	}

	commands := info.Build.Commands
	b.writeCommands(w, commands)

	return nil
}

func (b *AbstractShell) cacheArchiver(w ShellWriter, options *archivingOptions, info common.ShellScriptInfo) {
	if options == nil {
		return
	}

	// Skip archiving if no cache is defined
	cacheKey, cacheFile := b.cacheFile(info.Build, options.Key)
	if cacheKey == "" {
		return
	}

	args := []string{
		"cache-archiver",
		"--file", cacheFile,
	}

	// Create list of files to archive
	archiverArgs := options.CommandArguments()
	if len(archiverArgs) == 0 {
		// Skip creating archive
		return
	}
	args = append(args, archiverArgs...)

	// Generate cache upload address
	if url := getCacheUploadURL(info.Build, cacheKey); url != nil {
		args = append(args, "--url", url.String())
	}

	b.guardRunnerCommand(w, info.RunnerCommand, "Creating cache", func() {
		// Execute archive command
		w.Notice("Creating cache %s...", cacheKey)
		w.Command(info.RunnerCommand, args...)
	})
}

func (b *AbstractShell) uploadArtifacts(w ShellWriter, options *archivingOptions, info common.ShellScriptInfo) {
	if options == nil {
		return
	}
	if info.Build.Runner.URL == "" {
		return
	}

	args := []string{
		"artifacts-uploader",
		"--url",
		info.Build.Runner.URL,
		"--token",
		info.Build.Token,
		"--id",
		strconv.Itoa(info.Build.ID),
	}

	// Create list of files to archive
	archiverArgs := options.CommandArguments()
	if len(archiverArgs) == 0 {
		// Skip creating archive
		return
	}
	args = append(args, archiverArgs...)

	// Get artifacts:name
	if name, ok := info.Build.Options.GetString("artifacts", "name"); ok && name != "" {
		args = append(args, "--name", name)
	}

	// Get artifacts:expire_in
	if expireIn, ok := info.Build.Options.GetString("artifacts", "expire_in"); ok && expireIn != "" {
		args = append(args, "--expire-in", expireIn)
	}

	b.guardRunnerCommand(w, info.RunnerCommand, "Uploading artifacts", func() {
		w.Notice("Uploading artifacts...")
		w.Command(info.RunnerCommand, args...)
	})
}

func (b *AbstractShell) writeAfterScript(w ShellWriter, info common.ShellScriptInfo) error {
	shellOptions := struct {
		AfterScript []string `json:"after_script"`
	}{}
	err := info.Build.Options.Decode(&shellOptions)
	if err != nil {
		return err
	}

	if len(shellOptions.AfterScript) == 0 {
		return nil
	}

	b.writeExports(w, info)
	b.writeCdBuildDir(w, info)

	w.Notice("Running after script...")

	for _, command := range shellOptions.AfterScript {
		command = strings.TrimSpace(command)
		if command != "" {
			w.Notice("$ %s", command)
		} else {
			w.EmptyLine()
		}
		w.Line(command)
		w.CheckForErrors()
	}

	return nil
}

func (b *AbstractShell) writeArchiveCacheScript(w ShellWriter, info common.ShellScriptInfo) (err error) {
	// Parse options
	var options shellOptions
	err = info.Build.Options.Decode(&options)
	if err != nil {
		return
	}

	b.writeExports(w, info)
	b.writeCdBuildDir(w, info)
	b.writeTLSCAInfo(w, info.Build, "CI_SERVER_TLS_CA_FILE")

	// Find cached files and archive them
	b.cacheArchiver(w, options.Cache, info)
	return
}

func (b *AbstractShell) writeUploadArtifactsScript(w ShellWriter, info common.ShellScriptInfo) (err error) {
	// Parse options
	var options shellOptions
	err = info.Build.Options.Decode(&options)
	if err != nil {
		return
	}

	b.writeExports(w, info)
	b.writeCdBuildDir(w, info)
	b.writeTLSCAInfo(w, info.Build, "CI_SERVER_TLS_CA_FILE")

	// Upload artifacts
	b.uploadArtifacts(w, options.Artifacts, info)
	return
}

func (b *AbstractShell) writeScript(w ShellWriter, scriptType common.ShellScriptType, info common.ShellScriptInfo) (err error) {
	switch scriptType {
	case common.ShellPrepareScript:
		return b.writePrepareScript(w, info)

	case common.ShellBuildScript:
		return b.writeBuildScript(w, info)

	case common.ShellAfterScript:
		return b.writeAfterScript(w, info)

	case common.ShellArchiveCache:
		return b.writeArchiveCacheScript(w, info)

	case common.ShellUploadArtifacts:
		return b.writeUploadArtifactsScript(w, info)

	default:
		return errors.New("Not supported script type: " + string(scriptType))
	}
}
<|endoftext|>


<commit_before>package device

import (
	"fmt"
	"path/filepath"
	"strings"
	"time"

	log "github.com/sirupsen/logrus"
	"github.com/xshellinc/esp-flasher/serialport"
	"github.com/xshellinc/go-serial"
	"github.com/xshellinc/iotit/device/config"
	"github.com/xshellinc/iotit/workstation"
	"github.com/xshellinc/tools/dialogs"
	"github.com/xshellinc/tools/lib/help"
)

const (
	COLIBRI = "colibri"
)

var serialPort serial.Serial

const portSelectionTries = 3

// toradex colibri imx6 device
type colibri struct {
	*flasher
	Port string
	Disk string
}

// Prepare overrides flasher Prepare method with port initialization
func (d *colibri) Prepare() error {
	// start VM, upload image and extract it
	d.flasher.Prepare()
	log.WithField("device", COLIBRI).Debug("Prepare")
	// install toradex flasher dependencies
	d.installTools()

	d.getPort()
	return nil
}

func (d *colibri) getPort() error {
	if d.Port != "" {
		return nil
	}
	var perr error
	for attempt := 0; attempt < portSelectionTries; attempt++ {
		if attempt > 0 && !dialogs.YesNoDialog("No ports found. Reconnect your device and try again. Ready?") {
			return perr
		}
		fmt.Println("[+] Enumerating serial ports...")
		port, err := serialport.GetPort("auto")
		if err != nil {
			perr = err
			continue
		}
		d.Port = port
		break
	}
	if d.Port == "" {
		return perr
	}
	fmt.Println("[+] Using ", dialogs.PrintColored(d.Port))
	return nil
}

// Configure overrides flasher Configure() method with custom image configuration
func (d *colibri) Configure() error {
	log.WithField("device", COLIBRI).Debug("Configure")
	fmt.Println("[+] Configuring...")

	job := help.NewBackgroundJob()

	go func() {
		defer job.Close()

		if err := d.configureImage(); err != nil {
			job.Error(err)
			return
		}
	}()

	if err := help.WaitJobAndSpin("Preparing", job); err != nil {
		return err
	}

	fmt.Println("[+] Image ready")
	return nil
}

// Flash configures and flashes image
func (d *colibri) Flash() error {
	if err := d.Prepare(); err != nil {
		return err
	}

	if err := d.Configure(); err != nil {
		return err
	}

	if err := d.Write(); err != nil {
		return err
	}

	return d.Done()
}

func (d *colibri) installTools() error {
	if err := d.exec("apk add dosfstools parted sudo e2fsprogs-extra coreutils libattr"); err != nil {
		return err
	}
	return nil
}

func (d *colibri) configureImage() error {
	d.img = "colibri_image.zip"

	log.Debug("Creating tmp folder")
	command := fmt.Sprintf("mkdir -p %s", config.MountDir)
	if err := d.exec(command); err != nil {
		return err
	}
	log.Debug("Running update.sh")
	command = fmt.Sprintf("cd %s && ./update.sh -o %s", help.AddPathSuffix("unix", config.TmpDir, d.folder), config.MountDir)
	if err := d.exec(command); err != nil {
		return err
	}
	command = fmt.Sprintf("cd %s && zip -0 -r ../%s *", config.MountDir, d.img)
	if err := d.exec(command); err != nil {
		return err
	}
	return nil
}

// Write - writes image to SD card
func (d *colibri) Write() error {
	if d.img != "" {
		log.WithField("img", d.img).Debug("Downloading image from vbox")

		job := help.NewBackgroundJob()
		go func() {
			defer job.Close()
			if err := d.conf.SSH.ScpFrom(help.AddPathSuffix("unix", config.TmpDir, d.img), filepath.Join(help.GetTempDir(), d.img)); err != nil {
				job.Error(err)
			}
		}()

		if err := help.WaitJobAndSpin("Copying files", job); err != nil {
			log.Error(err)
			return err
		}
	} else {
		d.img = "colibri_image.zip"
	}

	flash_only := d.flasher.CLI.Bool("flash")
	if !flash_only {
		w := workstation.NewWorkStation(d.Disk)
		img := filepath.Join(help.GetTempDir(), d.img)

		log.WithField("img", img).Debug("Writing image to disk")

		if job, err := w.CopyToDisk(img); err != nil {
			return err
		} else if job != nil {
			if err := help.WaitJobAndSpin("Flashing", job); err != nil {
				return err
			}
		}
		time.Sleep(time.Second * 2)
		w.Unmount()
		fmt.Println("[+] SD card prepared")
	}
	if d.Port == "" {
		if err := d.getPort(); err != nil {
			log.Error(err)
			return err
		}
	}
	commonOpts := serial.OpenOptions{
		BaudRate:              115200,
		DataBits:              8,
		ParityMode:            serial.PARITY_NONE,
		StopBits:              1,
		InterCharacterTimeout: 200.0,
		PortName:              d.Port,
	}
	var err error
	serialPort, err = serial.Open(commonOpts)
	if err != nil {
		return err
	}

	defer serialPort.Close()

	if err := d.runUpdate(); err != nil {
		return err
	}

	return nil
}

func readResponse() string {
	data := make([]byte, 10000)
	if _, err := serialPort.Read(data); err != nil {
		log.Error(err)
	}
	line := strings.TrimSpace(string(data))
	log.WithField("data", line).Debug("Response")
	return line
}

func rebootBoard() error {
	serialPort.SetReadTimeout(time.Second * 5)
	if _, err := serialPort.Write([]byte("\r\n")); err != nil {
		log.Error(err)
		serialPort.Write([]byte("\r\n"))
	}
	for {
		line := readResponse()
		if line != "" {
			if strings.Contains(line, "imx6 login:") {
				if _, err := serialPort.Write([]byte("root\r\n")); err != nil {
					log.Error(err)
				}
				readResponse() //command echo
				readResponse() //response
				time.Sleep(time.Second * 2)
				if _, err := serialPort.Write([]byte("reboot\r\n")); err != nil {
					log.Error(err)
				}
				readResponse() //command echo
				readResponse() //response
				time.Sleep(time.Second * 2)
				return nil
			}
		}
	}
	return nil
}

func bootInRecovery() *help.BackgroundJob {
	job := help.NewBackgroundJob()
	go func() {
		defer job.Close()
		for {
			if _, err := serialPort.Write([]byte(" ")); err != nil {
				job.Error(err)
			}
			data := make([]byte, 10000)
			if _, err := serialPort.Read(data); err == nil {
				log.WithField("data", string(data)).Debug("Response")
				if strings.Contains(string(data), "Colibri iMX6 #") || strings.Contains(string(data), "iMX6 #") {
					break
				}
			}
		}
	}()
	return job
}

func (d *colibri) runUpdate() error {
	for !dialogs.YesNoDialog("Please insert prepared SD card into your Colibri iMX6 board. Type yes once ready.") {
	}
	message := "Now reset or power up the board"
	if !dialogs.YesNoDialog("Do you have a reset button on your carrier board?") {
		fmt.Println("[+] Trying to reboot the board")
		rebootBoard()
		message = "Booting in recovery"
	}
	job := bootInRecovery()
	if err := help.WaitJobAndSpin(message, job); err != nil {
		log.Error(err)
		return err
	}
	fmt.Println("[+] Flashing to eMMC")
	if _, err := serialPort.Write([]byte("run setupdate\r\n")); err != nil {
		log.Error(err)
	}

	readResponse()
	time.Sleep(time.Second * 1)
	readResponse()
	time.Sleep(time.Second * 2)
	if n, err := serialPort.Write([]byte("run update\r\n")); err != nil {
		log.Error(err)
	} else {
		log.Debug("Written:", n)
	}
	readResponse()
	for {
		serialPort.SetReadTimeout(time.Second * 5)
		line := readResponse()
		fmt.Print(line)
		if strings.Contains(line, "resetting") {
			fmt.Println("[+] Done! Rebooting the board.")
			return nil
		}
	}

	return nil
}

func (d *colibri) exec(command string) error {
	log.Debug(command)
	if out, eut, err := d.conf.SSH.Run(command); err != nil {
		log.Error("[-] Error executing: ", command, eut)
		return err
	} else if strings.TrimSpace(out) != "" {
		log.Debug(strings.TrimSpace(out))
	}
	return nil
}
<commit_msg>Remove port querying before it needed<commit_after>package device

import (
	"fmt"
	"path/filepath"
	"strings"
	"time"

	log "github.com/sirupsen/logrus"
	"github.com/xshellinc/esp-flasher/serialport"
	"github.com/xshellinc/go-serial"
	"github.com/xshellinc/iotit/device/config"
	"github.com/xshellinc/iotit/workstation"
	"github.com/xshellinc/tools/dialogs"
	"github.com/xshellinc/tools/lib/help"
)

const (
	COLIBRI = "colibri"
)

var serialPort serial.Serial

const portSelectionTries = 3

// toradex colibri imx6 device
type colibri struct {
	*flasher
	Port string
	Disk string
}

// Prepare overrides flasher Prepare method with port initialization
func (d *colibri) Prepare() error {
	// start VM, upload image and extract it
	d.flasher.Prepare()
	log.WithField("device", COLIBRI).Debug("Prepare")
	// install toradex flasher dependencies
	d.installTools()
	return nil
}

func (d *colibri) getPort() error {
	if d.Port != "" {
		return nil
	}
	var perr error
	for attempt := 0; attempt < portSelectionTries; attempt++ {
		if attempt > 0 && !dialogs.YesNoDialog("No ports found. Reconnect your device and try again. Ready?") {
			return perr
		}
		fmt.Println("[+] Enumerating serial ports...")
		port, err := serialport.GetPort("auto")
		if err != nil {
			perr = err
			continue
		}
		d.Port = port
		break
	}
	if d.Port == "" {
		return perr
	}
	fmt.Println("[+] Using ", dialogs.PrintColored(d.Port))
	return nil
}

// Configure overrides flasher Configure() method with custom image configuration
func (d *colibri) Configure() error {
	log.WithField("device", COLIBRI).Debug("Configure")
	fmt.Println("[+] Configuring...")

	job := help.NewBackgroundJob()

	go func() {
		defer job.Close()

		if err := d.configureImage(); err != nil {
			job.Error(err)
			return
		}
	}()

	if err := help.WaitJobAndSpin("Preparing", job); err != nil {
		return err
	}

	fmt.Println("[+] Image ready")
	return nil
}

// Flash configures and flashes image
func (d *colibri) Flash() error {
	if err := d.Prepare(); err != nil {
		return err
	}

	if err := d.Configure(); err != nil {
		return err
	}

	if err := d.Write(); err != nil {
		return err
	}

	return d.Done()
}

func (d *colibri) installTools() error {
	if err := d.exec("apk add dosfstools parted sudo e2fsprogs-extra coreutils libattr"); err != nil {
		return err
	}
	return nil
}

func (d *colibri) configureImage() error {
	d.img = "colibri_image.zip"

	log.Debug("Creating tmp folder")
	command := fmt.Sprintf("mkdir -p %s", config.MountDir)
	if err := d.exec(command); err != nil {
		return err
	}
	log.Debug("Running update.sh")
	command = fmt.Sprintf("cd %s && ./update.sh -o %s", help.AddPathSuffix("unix", config.TmpDir, d.folder), config.MountDir)
	if err := d.exec(command); err != nil {
		return err
	}
	command = fmt.Sprintf("cd %s && zip -0 -r ../%s *", config.MountDir, d.img)
	if err := d.exec(command); err != nil {
		return err
	}
	return nil
}

// Write - writes image to SD card
func (d *colibri) Write() error {
	if d.img != "" {
		log.WithField("img", d.img).Debug("Downloading image from vbox")

		job := help.NewBackgroundJob()
		go func() {
			defer job.Close()
			if err := d.conf.SSH.ScpFrom(help.AddPathSuffix("unix", config.TmpDir, d.img), filepath.Join(help.GetTempDir(), d.img)); err != nil {
				job.Error(err)
			}
		}()

		if err := help.WaitJobAndSpin("Copying files", job); err != nil {
			log.Error(err)
			return err
		}
	} else {
		d.img = "colibri_image.zip"
	}

	flash_only := d.flasher.CLI.Bool("flash")
	if !flash_only {
		w := workstation.NewWorkStation(d.Disk)
		img := filepath.Join(help.GetTempDir(), d.img)

		log.WithField("img", img).Debug("Writing image to disk")

		if job, err := w.CopyToDisk(img); err != nil {
			return err
		} else if job != nil {
			if err := help.WaitJobAndSpin("Writing image files to SD card", job); err != nil {
				return err
			}
		}
		time.Sleep(time.Second * 2)
		w.Unmount()
		fmt.Println("[+] SD card prepared")
	}
	if d.Port == "" {
		if err := d.getPort(); err != nil {
			log.Error(err)
			return err
		}
	}
	commonOpts := serial.OpenOptions{
		BaudRate:              115200,
		DataBits:              8,
		ParityMode:            serial.PARITY_NONE,
		StopBits:              1,
		InterCharacterTimeout: 200.0,
		PortName:              d.Port,
	}
	var err error
	serialPort, err = serial.Open(commonOpts)
	if err != nil {
		return err
	}

	defer serialPort.Close()

	if err := d.runUpdate(); err != nil {
		return err
	}

	return nil
}

func readResponse() string {
	data := make([]byte, 10000)
	if _, err := serialPort.Read(data); err != nil {
		log.Error(err)
	}
	line := strings.TrimSpace(string(data))
	log.WithField("data", line).Debug("Response")
	return line
}

func rebootBoard() error {
	serialPort.SetReadTimeout(time.Second * 5)
	if _, err := serialPort.Write([]byte("\r\n")); err != nil {
		log.Error(err)
		serialPort.Write([]byte("\r\n"))
	}
	for {
		line := readResponse()
		if line != "" {
			if strings.Contains(line, "imx6 login:") {
				if _, err := serialPort.Write([]byte("root\r\n")); err != nil {
					log.Error(err)
				}
				readResponse() //command echo
				readResponse() //response
				time.Sleep(time.Second * 2)
				if _, err := serialPort.Write([]byte("reboot\r\n")); err != nil {
					log.Error(err)
				}
				readResponse() //command echo
				readResponse() //response
				time.Sleep(time.Second * 2)
				return nil
			}
		}
	}
	return nil
}

func bootInRecovery() *help.BackgroundJob {
	job := help.NewBackgroundJob()
	go func() {
		defer job.Close()
		for {
			if _, err := serialPort.Write([]byte(" ")); err != nil {
				job.Error(err)
			}
			data := make([]byte, 10000)
			if _, err := serialPort.Read(data); err == nil {
				log.WithField("data", string(data)).Debug("Response")
				if strings.Contains(string(data), "Colibri iMX6 #") || strings.Contains(string(data), "iMX6 #") {
					break
				}
			}
		}
	}()
	return job
}

func (d *colibri) runUpdate() error {
	for !dialogs.YesNoDialog("Please insert prepared SD card into your Colibri iMX6 board. Type yes once ready.") {
	}
	message := "Now reset or power up the board"
	if !dialogs.YesNoDialog("Do you have a reset button on your carrier board?") {
		fmt.Println("[+] Trying to reboot the board")
		rebootBoard()
		message = "Booting in recovery"
	}
	job := bootInRecovery()
	if err := help.WaitJobAndSpin(message, job); err != nil {
		log.Error(err)
		return err
	}
	fmt.Println("[+] Flashing to eMMC")
	if _, err := serialPort.Write([]byte("run setupdate\r\n")); err != nil {
		log.Error(err)
	}

	readResponse()
	time.Sleep(time.Second * 1)
	readResponse()
	time.Sleep(time.Second * 2)
	if n, err := serialPort.Write([]byte("run update\r\n")); err != nil {
		log.Error(err)
	} else {
		log.Debug("Written:", n)
	}
	readResponse()
	for {
		serialPort.SetReadTimeout(time.Second * 5)
		line := readResponse()
		fmt.Print(line)
		if strings.Contains(line, "resetting") {
			fmt.Println("[+] Done! Rebooting the board.")
			return nil
		}
	}

	return nil
}

func (d *colibri) exec(command string) error {
	log.Debug(command)
	if out, eut, err := d.conf.SSH.Run(command); err != nil {
		log.Error("[-] Error executing: ", command, eut)
		return err
	} else if strings.TrimSpace(out) != "" {
		log.Debug(strings.TrimSpace(out))
	}
	return nil
}
<|endoftext|>