<commit_before>package client

import (
	"fmt"
	"strings"
	"sync"
	"time"

	"gopkg.in/jcmturner/gokrb5.v6/iana/nametype"
	"gopkg.in/jcmturner/gokrb5.v6/krberror"
	"gopkg.in/jcmturner/gokrb5.v6/messages"
	"gopkg.in/jcmturner/gokrb5.v6/types"
)

// sessions hold TGTs and are keyed on the realm name
type sessions struct {
	Entries map[string]*session
	mux     sync.RWMutex
}

// destroy erases all sessions
func (s *sessions) destroy() {
	s.mux.Lock()
	defer s.mux.Unlock()
	for k, e := range s.Entries {
		e.destroy()
		delete(s.Entries, k)
	}
}

// update replaces a session with the one provided or adds it as a new one
func (s *sessions) update(sess *session) {
	s.mux.Lock()
	defer s.mux.Unlock()
	// if a session already exists for this, cancel its auto renew.
	if i, ok := s.Entries[sess.realm]; ok {
		if i != sess {
			// Session in the sessions cache is not the same as one provided.
			// Cancel the one in the cache and add this one.
			i.mux.Lock()
			defer i.mux.Unlock()
			i.cancel <- true
			s.Entries[sess.realm] = sess
			return
		}
	}
	// No session for this realm was found so just add it
	s.Entries[sess.realm] = sess
}

// get returns the session for the realm specified
func (s *sessions) get(realm string) (*session, bool) {
	s.mux.RLock()
	defer s.mux.RUnlock()
	sess, ok := s.Entries[realm]
	return sess, ok
}

// session holds the TGT details for a realm
type session struct {
	realm                string
	authTime             time.Time
	endTime              time.Time
	renewTill            time.Time
	tgt                  messages.Ticket
	sessionKey           types.EncryptionKey
	sessionKeyExpiration time.Time
	cancel               chan bool
	mux                  sync.RWMutex
}

// AddSession adds a session for a realm with a TGT to the client's session cache.
// A goroutine is started to automatically renew the TGT before expiry.
func (cl *Client) AddSession(tgt messages.Ticket, dep messages.EncKDCRepPart) {
	if strings.ToLower(tgt.SName.NameString[0]) != "krbtgt" {
		// Not a TGT
		return
	}
	s := &session{
		realm:                tgt.SName.NameString[len(tgt.SName.NameString)-1],
		authTime:             dep.AuthTime,
		endTime:              dep.EndTime,
		renewTill:            dep.RenewTill,
		tgt:                  tgt,
		sessionKey:           dep.Key,
		sessionKeyExpiration: dep.KeyExpiration,
	}
	cl.sessions.update(s)
	cl.enableAutoSessionRenewal(s)
}

// update overwrites the session details with those from the TGT and decrypted encPart
func (s *session) update(tgt messages.Ticket, dep messages.EncKDCRepPart) {
	s.mux.Lock()
	defer s.mux.Unlock()
	s.authTime = dep.AuthTime
	s.endTime = dep.EndTime
	s.renewTill = dep.RenewTill
	s.tgt = tgt
	s.sessionKey = dep.Key
	s.sessionKeyExpiration = dep.KeyExpiration
}

// destroy will cancel any auto renewal of the session and set the expiration times to the current time
func (s *session) destroy() {
	s.mux.Lock()
	defer s.mux.Unlock()
	if s.cancel != nil {
		s.cancel <- true
	}
	s.endTime = time.Now().UTC()
	s.renewTill = s.endTime
	s.sessionKeyExpiration = s.endTime
}

// valid informs if the TGT is still within the valid time window
func (s *session) valid() bool {
	s.mux.RLock()
	defer s.mux.RUnlock()
	t := time.Now().UTC()
	if t.Before(s.endTime) && s.authTime.Before(t) {
		return true
	}
	return false
}

// tgtDetails is a thread safe way to get the session's realm, TGT and session key values
func (s *session) tgtDetails() (string, messages.Ticket, types.EncryptionKey) {
	s.mux.RLock()
	defer s.mux.RUnlock()
	return s.realm, s.tgt, s.sessionKey
}

// timeDetails is a thread safe way to get the session's validity time values
func (s *session) timeDetails() (string, time.Time, time.Time, time.Time, time.Time) {
	s.mux.RLock()
	defer s.mux.RUnlock()
	return s.realm, s.authTime, s.endTime, s.renewTill, s.sessionKeyExpiration
}

// enableAutoSessionRenewal turns on the automatic renewal for the client's TGT session.
func (cl *Client) enableAutoSessionRenewal(s *session) {
	var timer *time.Timer
	s.cancel = make(chan bool, 1)
	go func(s *session) {
		for {
			s.mux.RLock()
			w := (s.endTime.Sub(time.Now().UTC()) * 5) / 6
			s.mux.RUnlock()
			if w < 0 {
				return
			}
			timer = time.NewTimer(w)
			select {
			case <-timer.C:
				renewal, err := cl.refreshSession(s)
				if !renewal && err == nil {
					// end this goroutine as there will have been a new login and new auto renewal goroutine created.
					return
				}
			case <-s.cancel:
				// cancel has been called. Stop the timer and exit.
				timer.Stop()
				return
			}
		}
	}(s)
}

// renewTGT renews the client's TGT session.
func (cl *Client) renewTGT(s *session) error {
	realm, tgt, skey := s.tgtDetails()
	spn := types.PrincipalName{
		NameType:   nametype.KRB_NT_SRV_INST,
		NameString: []string{"krbtgt", realm},
	}
	_, tgsRep, err := cl.TGSExchange(spn, cl.Credentials.Realm, tgt, skey, true, 0)
	if err != nil {
		return krberror.Errorf(err, krberror.KRBMsgError, "error renewing TGT")
	}
	s.update(tgsRep.Ticket, tgsRep.DecryptedEncPart)
	cl.sessions.update(s)
	return nil
}

// refreshSession updates either through renewal or creating a new login.
// The boolean indicates if the update was a renewal.
func (cl *Client) refreshSession(s *session) (bool, error) {
	s.mux.RLock()
	realm := s.realm
	renewTill := s.renewTill
	s.mux.RUnlock()
	if time.Now().UTC().Before(renewTill) {
		err := cl.renewTGT(s)
		return true, err
	}
	err := cl.realmLogin(realm)
	return false, err
}

// ensureValidSession makes sure there is a valid session for the realm
func (cl *Client) ensureValidSession(realm string) error {
	s, ok := cl.sessions.get(realm)
	if ok {
		s.mux.RLock()
		defer s.mux.RUnlock()
		d := s.endTime.Sub(s.authTime) / 6
		if s.endTime.Sub(time.Now().UTC()) > d {
			return nil
		}
		_, err := cl.refreshSession(s)
		return err
	}
	return cl.realmLogin(realm)
}

// sessionTGTDetails is a thread safe way to get the TGT and session key values for a realm
func (cl *Client) sessionTGT(realm string) (tgt messages.Ticket, sessionKey types.EncryptionKey, err error) {
	err = cl.ensureValidSession(realm)
	if err != nil {
		return
	}
	s, ok := cl.sessions.get(realm)
	if !ok {
		err = fmt.Errorf("could not find TGT session for %s", realm)
		return
	}
	_, tgt, sessionKey = s.tgtDetails()
	return
}

func (cl *Client) sessionTimes(realm string) (authTime, endTime, renewTime, sessionExp time.Time, err error) {
	s, ok := cl.sessions.get(realm)
	if !ok {
		err = fmt.Errorf("could not find TGT session for %s", realm)
		return
	}
	_, authTime, endTime, renewTime, sessionExp = s.timeDetails()
	return
}

// spnRealm resolves the realm name of a service principal name
func (cl *Client) spnRealm(spn types.PrincipalName) string {
	return cl.Config.ResolveRealm(spn.NameString[len(spn.NameString)-1])
}
<commit_msg>race condition fix<commit_after>package client

import (
	"fmt"
	"strings"
	"sync"
	"time"

	"gopkg.in/jcmturner/gokrb5.v6/iana/nametype"
	"gopkg.in/jcmturner/gokrb5.v6/krberror"
	"gopkg.in/jcmturner/gokrb5.v6/messages"
	"gopkg.in/jcmturner/gokrb5.v6/types"
)

// sessions hold TGTs and are keyed on the realm name
type sessions struct {
	Entries map[string]*session
	mux     sync.RWMutex
}

// destroy erases all sessions
func (s *sessions) destroy() {
	s.mux.Lock()
	defer s.mux.Unlock()
	for k, e := range s.Entries {
		e.destroy()
		delete(s.Entries, k)
	}
}

// update replaces a session with the one provided or adds it as a new one
func (s *sessions) update(sess *session) {
	s.mux.Lock()
	defer s.mux.Unlock()
	// if a session already exists for this, cancel its auto renew.
	if i, ok := s.Entries[sess.realm]; ok {
		if i != sess {
			// Session in the sessions cache is not the same as one provided.
			// Cancel the one in the cache and add this one.
			i.mux.Lock()
			defer i.mux.Unlock()
			i.cancel <- true
			s.Entries[sess.realm] = sess
			return
		}
	}
	// No session for this realm was found so just add it
	s.Entries[sess.realm] = sess
}

// get returns the session for the realm specified
func (s *sessions) get(realm string) (*session, bool) {
	s.mux.RLock()
	defer s.mux.RUnlock()
	sess, ok := s.Entries[realm]
	return sess, ok
}

// session holds the TGT details for a realm
type session struct {
	realm                string
	authTime             time.Time
	endTime              time.Time
	renewTill            time.Time
	tgt                  messages.Ticket
	sessionKey           types.EncryptionKey
	sessionKeyExpiration time.Time
	cancel               chan bool
	mux                  sync.RWMutex
}

// AddSession adds a session for a realm with a TGT to the client's session cache.
// A goroutine is started to automatically renew the TGT before expiry.
func (cl *Client) AddSession(tgt messages.Ticket, dep messages.EncKDCRepPart) {
	if strings.ToLower(tgt.SName.NameString[0]) != "krbtgt" {
		// Not a TGT
		return
	}
	s := &session{
		realm:                tgt.SName.NameString[len(tgt.SName.NameString)-1],
		authTime:             dep.AuthTime,
		endTime:              dep.EndTime,
		renewTill:            dep.RenewTill,
		tgt:                  tgt,
		sessionKey:           dep.Key,
		sessionKeyExpiration: dep.KeyExpiration,
	}
	cl.sessions.update(s)
	cl.enableAutoSessionRenewal(s)
}

// update overwrites the session details with those from the TGT and decrypted encPart
func (s *session) update(tgt messages.Ticket, dep messages.EncKDCRepPart) {
	s.mux.Lock()
	defer s.mux.Unlock()
	s.authTime = dep.AuthTime
	s.endTime = dep.EndTime
	s.renewTill = dep.RenewTill
	s.tgt = tgt
	s.sessionKey = dep.Key
	s.sessionKeyExpiration = dep.KeyExpiration
}

// destroy will cancel any auto renewal of the session and set the expiration times to the current time
func (s *session) destroy() {
	s.mux.Lock()
	defer s.mux.Unlock()
	if s.cancel != nil {
		s.cancel <- true
	}
	s.endTime = time.Now().UTC()
	s.renewTill = s.endTime
	s.sessionKeyExpiration = s.endTime
}

// valid informs if the TGT is still within the valid time window
func (s *session) valid() bool {
	s.mux.RLock()
	defer s.mux.RUnlock()
	t := time.Now().UTC()
	if t.Before(s.endTime) && s.authTime.Before(t) {
		return true
	}
	return false
}

// tgtDetails is a thread safe way to get the session's realm, TGT and session key values
func (s *session) tgtDetails() (string, messages.Ticket, types.EncryptionKey) {
	s.mux.RLock()
	defer s.mux.RUnlock()
	return s.realm, s.tgt, s.sessionKey
}

// timeDetails is a thread safe way to get the session's validity time values
func (s *session) timeDetails() (string, time.Time, time.Time, time.Time, time.Time) {
	s.mux.RLock()
	defer s.mux.RUnlock()
	return s.realm, s.authTime, s.endTime, s.renewTill, s.sessionKeyExpiration
}

// enableAutoSessionRenewal turns on the automatic renewal for the client's TGT session.
func (cl *Client) enableAutoSessionRenewal(s *session) {
	var timer *time.Timer
	s.mux.Lock()
	s.cancel = make(chan bool, 1)
	s.mux.Unlock()
	go func(s *session) {
		for {
			s.mux.RLock()
			w := (s.endTime.Sub(time.Now().UTC()) * 5) / 6
			s.mux.RUnlock()
			if w < 0 {
				return
			}
			timer = time.NewTimer(w)
			select {
			case <-timer.C:
				renewal, err := cl.refreshSession(s)
				if !renewal && err == nil {
					// end this goroutine as there will have been a new login and new auto renewal goroutine created.
					return
				}
			case <-s.cancel:
				// cancel has been called. Stop the timer and exit.
				timer.Stop()
				return
			}
		}
	}(s)
}

// renewTGT renews the client's TGT session.
func (cl *Client) renewTGT(s *session) error {
	realm, tgt, skey := s.tgtDetails()
	spn := types.PrincipalName{
		NameType:   nametype.KRB_NT_SRV_INST,
		NameString: []string{"krbtgt", realm},
	}
	_, tgsRep, err := cl.TGSExchange(spn, cl.Credentials.Realm, tgt, skey, true, 0)
	if err != nil {
		return krberror.Errorf(err, krberror.KRBMsgError, "error renewing TGT")
	}
	s.update(tgsRep.Ticket, tgsRep.DecryptedEncPart)
	cl.sessions.update(s)
	return nil
}

// refreshSession updates either through renewal or creating a new login.
// The boolean indicates if the update was a renewal.
func (cl *Client) refreshSession(s *session) (bool, error) {
	s.mux.RLock()
	realm := s.realm
	renewTill := s.renewTill
	s.mux.RUnlock()
	if time.Now().UTC().Before(renewTill) {
		err := cl.renewTGT(s)
		return true, err
	}
	err := cl.realmLogin(realm)
	return false, err
}

// ensureValidSession makes sure there is a valid session for the realm
func (cl *Client) ensureValidSession(realm string) error {
	s, ok := cl.sessions.get(realm)
	if ok {
		s.mux.RLock()
		defer s.mux.RUnlock()
		d := s.endTime.Sub(s.authTime) / 6
		if s.endTime.Sub(time.Now().UTC()) > d {
			return nil
		}
		_, err := cl.refreshSession(s)
		return err
	}
	return cl.realmLogin(realm)
}

// sessionTGTDetails is a thread safe way to get the TGT and session key values for a realm
func (cl *Client) sessionTGT(realm string) (tgt messages.Ticket, sessionKey types.EncryptionKey, err error) {
	err = cl.ensureValidSession(realm)
	if err != nil {
		return
	}
	s, ok := cl.sessions.get(realm)
	if !ok {
		err = fmt.Errorf("could not find TGT session for %s", realm)
		return
	}
	_, tgt, sessionKey = s.tgtDetails()
	return
}

func (cl *Client) sessionTimes(realm string) (authTime, endTime, renewTime, sessionExp time.Time, err error) {
	s, ok := cl.sessions.get(realm)
	if !ok {
		err = fmt.Errorf("could not find TGT session for %s", realm)
		return
	}
	_, authTime, endTime, renewTime, sessionExp = s.timeDetails()
	return
}

// spnRealm resolves the realm name of a service principal name
func (cl *Client) spnRealm(spn types.PrincipalName) string {
	return cl.Config.ResolveRealm(spn.NameString[len(spn.NameString)-1])
}
<|endoftext|>


<commit_before>package client

import (
	"time"

	"github.com/hashicorp/nomad/client/config"
	consulApi "github.com/hashicorp/nomad/client/consul"
	"github.com/hashicorp/nomad/client/fingerprint"
	"github.com/hashicorp/nomad/command/agent/consul"
	"github.com/hashicorp/nomad/helper/testlog"
	"github.com/hashicorp/nomad/plugins/shared/catalog"
	"github.com/hashicorp/nomad/plugins/shared/singleton"
	"github.com/mitchellh/go-testing-interface"
)

// TestClient creates an in-memory client for testing purposes and returns a
// cleanup func to shutdown the client and remove the alloc and state dirs.
//
// There is no need to override the AllocDir or StateDir as they are randomized
// and removed in the returned cleanup function. If they are overridden in the
// callback then the caller still must run the returned cleanup func.
func TestClient(t testing.T, cb func(c *config.Config)) (*Client, func()) {
	conf, cleanup := config.TestClientConfig(t)

	// Tighten the fingerprinter timeouts (must be done in client package
	// to avoid circular dependencies)
	if conf.Options == nil {
		conf.Options = make(map[string]string)
	}
	conf.Options[fingerprint.TightenNetworkTimeoutsConfig] = "true"

	logger := testlog.HCLogger(t)
	conf.Logger = logger

	if cb != nil {
		cb(conf)
	}

	// Set the plugin loaders
	if conf.PluginLoader == nil {
		conf.PluginLoader = catalog.TestPluginLoaderWithOptions(t, "", conf.Options, nil)
		conf.PluginSingletonLoader = singleton.NewSingletonLoader(logger, conf.PluginLoader)
	}
	catalog := consul.NewMockCatalog(logger)
	mockService := consulApi.NewMockConsulServiceClient(t, logger)
	client, err := NewClient(conf, catalog, mockService)
	if err != nil {
		cleanup()
		t.Fatalf("err: %v", err)
	}
	return client, func() {
		ch := make(chan error, 1)

		go func() {
			defer close(ch)

			// Shutdown client
			err := client.Shutdown()
			if err != nil {
				t.Errorf("failed to shutdown client: %v", err)
			}

			// Call TestClientConfig cleanup
			cleanup()
		}()

		select {
		case <-ch:
			// all good
		case <-time.After(1 * time.Minute):
			t.Errorf("timed out cleaning up test client")
		}
	}
}
<commit_msg>tests: no need for buffer channel<commit_after>package client

import (
	"time"

	"github.com/hashicorp/nomad/client/config"
	consulApi "github.com/hashicorp/nomad/client/consul"
	"github.com/hashicorp/nomad/client/fingerprint"
	"github.com/hashicorp/nomad/command/agent/consul"
	"github.com/hashicorp/nomad/helper/testlog"
	"github.com/hashicorp/nomad/plugins/shared/catalog"
	"github.com/hashicorp/nomad/plugins/shared/singleton"
	"github.com/mitchellh/go-testing-interface"
)

// TestClient creates an in-memory client for testing purposes and returns a
// cleanup func to shutdown the client and remove the alloc and state dirs.
//
// There is no need to override the AllocDir or StateDir as they are randomized
// and removed in the returned cleanup function. If they are overridden in the
// callback then the caller still must run the returned cleanup func.
func TestClient(t testing.T, cb func(c *config.Config)) (*Client, func()) {
	conf, cleanup := config.TestClientConfig(t)

	// Tighten the fingerprinter timeouts (must be done in client package
	// to avoid circular dependencies)
	if conf.Options == nil {
		conf.Options = make(map[string]string)
	}
	conf.Options[fingerprint.TightenNetworkTimeoutsConfig] = "true"

	logger := testlog.HCLogger(t)
	conf.Logger = logger

	if cb != nil {
		cb(conf)
	}

	// Set the plugin loaders
	if conf.PluginLoader == nil {
		conf.PluginLoader = catalog.TestPluginLoaderWithOptions(t, "", conf.Options, nil)
		conf.PluginSingletonLoader = singleton.NewSingletonLoader(logger, conf.PluginLoader)
	}
	catalog := consul.NewMockCatalog(logger)
	mockService := consulApi.NewMockConsulServiceClient(t, logger)
	client, err := NewClient(conf, catalog, mockService)
	if err != nil {
		cleanup()
		t.Fatalf("err: %v", err)
	}
	return client, func() {
		ch := make(chan error)

		go func() {
			defer close(ch)

			// Shutdown client
			err := client.Shutdown()
			if err != nil {
				t.Errorf("failed to shutdown client: %v", err)
			}

			// Call TestClientConfig cleanup
			cleanup()
		}()

		select {
		case <-ch:
			// all good
		case <-time.After(1 * time.Minute):
			t.Errorf("timed out cleaning up test client")
		}
	}
}
<|endoftext|>


<commit_before>package helpers

import (
	"errors"
	"koding/db/mongodb"
	"koding/db/mongodb/modelhelper"
	"koding/tools/logger"

	"labix.org/v2/mgo"
)

// iterOptions holds the related config paramters for Iter operation
type iterOptions struct {
	// Starting offset
	Skip int

	// Ending point, iter count
	Limit int

	// Filter for limiting the result set
	Filter modelhelper.Selector

	// Iteration collection
	CollectionName string

	// Iteration function, all results will be passed to this function
	F func(result interface{})

	// Sometimes iteration can timeout, this is retry count
	RetryCount int

	// Data object itself for marshalling the result
	DataType interface{}

	// logger for iteration
	Log logger.Log
}

// NewIterOptions Sets the default values for iterOptions
func NewIterOptions() *iterOptions {
	return &iterOptions{
		Skip:       0,
		Limit:      1000,
		Filter:     modelhelper.Selector{},
		RetryCount: 50,
		Log:        logger.New("Iter"),
	}
}

// Iter accepts mongo and iterOptions and runs the query
func Iter(mongo *mongodb.MongoDB, iterOptions *iterOptions) error {
	if iterOptions.CollectionName == "" {
		return errors.New("Collection name is not set")
	}

	if iterOptions.F == nil {
		return errors.New("Iteration function is not given")
	}
	if iterOptions.DataType == nil {
		return errors.New("Datatype is not given")
	}

	return mongo.Run("jAccounts", createQuery(iterOptions))
}

// createQuery creates mongo query for iteration
func createQuery(iterOptions *iterOptions) func(coll *mgo.Collection) error {
	return func(coll *mgo.Collection) error {
		// find the total count
		query := coll.Find(iterOptions.Filter)
		totalCount, err := query.Count()
		if err != nil {
			iterOptions.Log.Error("While getting count, exiting: %v", err)
			return err
		}
		iterOptions.Log.Info("Totaly we have %v items for operation", totalCount)

		skip := iterOptions.Skip
		// this is a starting point
		index := skip
		// this is the item count to be processed
		limit := iterOptions.Limit
		// this will be the ending point
		count := index + limit

		iteration := 0
		for {
			// if we reach to the end of the all collection, exit
			if index >= totalCount {
				iterOptions.Log.Info("All items are processed, exiting")
				break
			}

			// this is the max re-iterating count
			if iteration == iterOptions.RetryCount {
				break
			}

			// if we processed all items then exit
			if index == count {
				break
			}

			iter := query.Skip(index).Limit(count - index).Iter()

			for iter.Next(iterOptions.DataType) {
				iterOptions.F(iterOptions.DataType)
				index++
				iterOptions.Log.Debug("Index: %v", index)
			}

			if err := iter.Close(); err != nil {
				iterOptions.Log.Error("Iteration failed: %v", err)
			}

			if iter.Timeout() {
				continue
			}

			iterOptions.Log.Info("iter existed, starting over from %v  -- %v  item(s) are processsed on this iter", index+1, index-skip)
			iteration++
		}

		if iteration == iterOptions.RetryCount {
			iterOptions.Log.Info("Max iteration count %v reached, exiting", iteration)
		}
		iterOptions.Log.Info("Deleted %v guest accounts on this process", index-skip)

		return nil
	}
}
<commit_msg>Iter: fix hardcoded collection name<commit_after>package helpers

import (
	"errors"
	"koding/db/mongodb"
	"koding/db/mongodb/modelhelper"
	"koding/tools/logger"

	"labix.org/v2/mgo"
)

// iterOptions holds the related config paramters for Iter operation
type iterOptions struct {
	// Starting offset
	Skip int

	// Ending point, iter count
	Limit int

	// Filter for limiting the result set
	Filter modelhelper.Selector

	// Iteration collection
	CollectionName string

	// Iteration function, all results will be passed to this function
	F func(result interface{})

	// Sometimes iteration can timeout, this is retry count
	RetryCount int

	// Data object itself for marshalling the result
	DataType interface{}

	// logger for iteration
	Log logger.Log
}

// NewIterOptions Sets the default values for iterOptions
func NewIterOptions() *iterOptions {
	return &iterOptions{
		Skip:       0,
		Limit:      1000,
		Filter:     modelhelper.Selector{},
		RetryCount: 50,
		Log:        logger.New("Iter"),
	}
}

// Iter accepts mongo and iterOptions and runs the query
func Iter(mongo *mongodb.MongoDB, iterOptions *iterOptions) error {
	if iterOptions.CollectionName == "" {
		return errors.New("Collection name is not set")
	}

	if iterOptions.F == nil {
		return errors.New("Iteration function is not given")
	}
	if iterOptions.DataType == nil {
		return errors.New("Datatype is not given")
	}

	return mongo.Run(iterOptions.CollectionName, createQuery(iterOptions))
}

// createQuery creates mongo query for iteration
func createQuery(iterOptions *iterOptions) func(coll *mgo.Collection) error {
	return func(coll *mgo.Collection) error {
		// find the total count
		query := coll.Find(iterOptions.Filter)
		totalCount, err := query.Count()
		if err != nil {
			iterOptions.Log.Error("While getting count, exiting: %v", err)
			return err
		}
		iterOptions.Log.Info("Totaly we have %v items for operation", totalCount)

		skip := iterOptions.Skip
		// this is a starting point
		index := skip
		// this is the item count to be processed
		limit := iterOptions.Limit
		// this will be the ending point
		count := index + limit

		iteration := 0
		for {
			// if we reach to the end of the all collection, exit
			if index >= totalCount {
				iterOptions.Log.Info("All items are processed, exiting")
				break
			}

			// this is the max re-iterating count
			if iteration == iterOptions.RetryCount {
				break
			}

			// if we processed all items then exit
			if index == count {
				break
			}

			iter := query.Skip(index).Limit(count - index).Iter()

			for iter.Next(iterOptions.DataType) {
				iterOptions.F(iterOptions.DataType)
				index++
				iterOptions.Log.Debug("Index: %v", index)
			}

			if err := iter.Close(); err != nil {
				iterOptions.Log.Error("Iteration failed: %v", err)
			}

			if iter.Timeout() {
				continue
			}

			iterOptions.Log.Info("iter existed, starting over from %v  -- %v  item(s) are processsed on this iter", index+1, index-skip)
			iteration++
		}

		if iteration == iterOptions.RetryCount {
			iterOptions.Log.Info("Max iteration count %v reached, exiting", iteration)
		}
		iterOptions.Log.Info("Deleted %v items on this process", index-skip)

		return nil
	}
}
<|endoftext|>


<commit_before>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package mysqlctl

import (
	"testing"
)

func TestStartShutdown(t *testing.T) {
	mycnf0 := NewMycnf(0, 3700, VtReplParams{})
	mycnf1 := NewMycnf(1, 3701, VtReplParams{})
	tablet0 := NewMysqld(mycnf0, DefaultDbaParams, DefaultReplParams)
	tablet1 := NewMysqld(mycnf1, DefaultDbaParams, DefaultReplParams)
	var err error

	err = Init(tablet0)
	if err != nil {
		t.Fatalf("Init(0) err: %v", err)
	}

	err = Init(tablet1)
	if err != nil {
		t.Fatalf("Init(1) err: %v", err)
	}

	err = Shutdown(tablet0, true)
	if err != nil {
		t.Fatalf("Shutdown() err: %v", err)
	}

	err = Start(tablet0)
	if err != nil {
		t.Fatalf("Start() err: %v", err)
	}

	err = Teardown(tablet0, false)
	if err != nil {
		t.Fatalf("Teardown(0) err: %v", err)
	}
	err = Teardown(tablet1, false)
	if err != nil {
		t.Fatalf("Teardown(1) err: %v", err)
	}
}
<commit_msg>Fixing this test I broke earlier.<commit_after>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package mysqlctl

import (
	"testing"
)

func TestStartShutdown(t *testing.T) {
	mycnf0 := NewMycnf(0, 3700, VtReplParams{})
	mycnf1 := NewMycnf(1, 3701, VtReplParams{})
	tablet0 := NewMysqld(mycnf0, DefaultDbaParams, DefaultReplParams)
	tablet1 := NewMysqld(mycnf1, DefaultDbaParams, DefaultReplParams)
	var err error

	err = Init(tablet0)
	if err != nil {
		t.Fatalf("Init(0) err: %v", err)
	}

	err = Init(tablet1)
	if err != nil {
		t.Fatalf("Init(1) err: %v", err)
	}

	err = Shutdown(tablet0, true, MysqlWaitTime)
	if err != nil {
		t.Fatalf("Shutdown() err: %v", err)
	}

	err = Start(tablet0, MysqlWaitTime)
	if err != nil {
		t.Fatalf("Start() err: %v", err)
	}

	err = Teardown(tablet0, false)
	if err != nil {
		t.Fatalf("Teardown(0) err: %v", err)
	}
	err = Teardown(tablet1, false)
	if err != nil {
		t.Fatalf("Teardown(1) err: %v", err)
	}
}
<|endoftext|>


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

package gitlab

import (
	"net/http"

	retryablehttp "github.com/hashicorp/go-retryablehttp"
)

// ClientOptionFunc can be used to customize a new GitLab API client.
type ClientOptionFunc func(*Client) error

// WithBaseURL sets the base URL for API requests to a custom endpoint.
func WithBaseURL(urlStr string) ClientOptionFunc {
	return func(c *Client) error {
		return c.setBaseURL(urlStr)
	}
}

// WithCustomBackoff can be used to configure a custom backoff policy.
func WithCustomBackoff(backoff retryablehttp.Backoff) ClientOptionFunc {
	return func(c *Client) error {
		c.client.Backoff = backoff
		return nil
	}
}

// WithCustomLimiter injects a custom rate limiter to the client.
func WithCustomLimiter(limiter RateLimiter) ClientOptionFunc {
	return func(c *Client) error {
		c.configureLimiterOnce.Do(func() {
			c.limiter = limiter
		})
		return nil
	}
}

// WithCustomRetry can be used to configure a custom retry policy.
func WithCustomRetry(checkRetry retryablehttp.CheckRetry) ClientOptionFunc {
	return func(c *Client) error {
		c.client.CheckRetry = checkRetry
		return nil
	}
}

// WithHTTPClient can be used to configure a custom HTTP client.
func WithHTTPClient(httpClient *http.Client) ClientOptionFunc {
	return func(c *Client) error {
		c.client.HTTPClient = httpClient
		return nil
	}
}

// WithoutRetries disables the default retry logic.
func WithoutRetries() ClientOptionFunc {
	return func(c *Client) error {
		c.disableRetries = true
		return nil
	}
}
<commit_msg>client_options: let user pass custom Logger and LeveledLogger<commit_after>//
// Copyright 2021, Sander van Harmelen
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

package gitlab

import (
	"net/http"

	retryablehttp "github.com/hashicorp/go-retryablehttp"
)

// ClientOptionFunc can be used to customize a new GitLab API client.
type ClientOptionFunc func(*Client) error

// WithBaseURL sets the base URL for API requests to a custom endpoint.
func WithBaseURL(urlStr string) ClientOptionFunc {
	return func(c *Client) error {
		return c.setBaseURL(urlStr)
	}
}

// WithCustomBackoff can be used to configure a custom backoff policy.
func WithCustomBackoff(backoff retryablehttp.Backoff) ClientOptionFunc {
	return func(c *Client) error {
		c.client.Backoff = backoff
		return nil
	}
}

// WithCustomLogger can be used to configure a custom retryablehttp leveled logger
func WithCustomLeveledLogger(leveledLogger retryablehttp.LeveledLogger) ClientOptionFunc {
	return func(c *Client) error {
		c.client.Logger = leveledLogger
		return nil
	}
}

// WithCustomLimiter injects a custom rate limiter to the client.
func WithCustomLimiter(limiter RateLimiter) ClientOptionFunc {
	return func(c *Client) error {
		c.configureLimiterOnce.Do(func() {
			c.limiter = limiter
		})
		return nil
	}
}

// WithCustomLogger can be used to configure a custom retryablehttp logger
func WithCustomLogger(logger retryablehttp.Logger) ClientOptionFunc {
	return func(c *Client) error {
		c.client.Logger = logger
		return nil
	}
}

// WithCustomRetry can be used to configure a custom retry policy.
func WithCustomRetry(checkRetry retryablehttp.CheckRetry) ClientOptionFunc {
	return func(c *Client) error {
		c.client.CheckRetry = checkRetry
		return nil
	}
}

// WithHTTPClient can be used to configure a custom HTTP client.
func WithHTTPClient(httpClient *http.Client) ClientOptionFunc {
	return func(c *Client) error {
		c.client.HTTPClient = httpClient
		return nil
	}
}

// WithoutRetries disables the default retry logic.
func WithoutRetries() ClientOptionFunc {
	return func(c *Client) error {
		c.disableRetries = true
		return nil
	}
}
<|endoftext|>


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

package couler

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"reflect"
	"strconv"
	"strings"

	"sqlflow.org/sqlflow/go/ir"
	pb "sqlflow.org/sqlflow/go/proto"
)

var defaultDockerImage = "sqlflow/sqlflow"
var workflowTTL = 24 * 3600
var envResource = "SQLFLOW_WORKFLOW_RESOURCES"

// Codegen generates Couler program
type Codegen struct{}

func fillMapIfValueNotEmpty(m map[string]string, key, value string) {
	if value != "" {
		m[key] = value
	}
}

func fillEnvFromSession(envs *map[string]string, session *pb.Session) {
	fillMapIfValueNotEmpty(*envs, "SQLFLOW_USER_TOKEN", session.Token)
	fillMapIfValueNotEmpty(*envs, "SQLFLOW_DATASOURCE", session.DbConnStr)
	fillMapIfValueNotEmpty(*envs, "SQLFLOW_USER_ID", session.UserId)
	fillMapIfValueNotEmpty(*envs, "SQLFLOW_HIVE_LOCATION", session.HiveLocation)
	fillMapIfValueNotEmpty(*envs, "SQLFLOW_HDFS_NAMENODE_ADDR", session.HdfsNamenodeAddr)
	fillMapIfValueNotEmpty(*envs, "SQLFLOW_HADOOP_USER", session.HdfsUser)
	fillMapIfValueNotEmpty(*envs, "SQLFLOW_HADOOP_PASS", session.HdfsUser)
	fillMapIfValueNotEmpty(*envs, "SQLFLOW_submitter", session.Submitter)
	fillMapIfValueNotEmpty(*envs, "SQLFLOW_WORKFLOW_STEP_LOG_FILE", os.Getenv("SQLFLOW_WORKFLOW_STEP_LOG_FILE"))
}

// GetStepEnvs returns a map of envs used for couler workflow.
func GetStepEnvs(session *pb.Session) (map[string]string, error) {
	envs := make(map[string]string)
	// fill step envs from the environment variables on sqlflowserver
	for _, env := range os.Environ() {
		pair := strings.SplitN(env, "=", 2)
		if len(pair) != 2 {
			return nil, fmt.Errorf("env: %s should format key=value", env)
		}
		// should not pass the workflow env into step
		if strings.HasPrefix(pair[0], "SQLFLOW_") && !strings.HasPrefix(pair[0], "SQLFLOW_WORKFLOW_") {
			envs[pair[0]] = pair[1]
		}
	}
	fillEnvFromSession(&envs, session)
	return envs, nil
}

// VerifyResources verifies the SQLFLOW_WORKFLOW_RESOURCES env to be valid.
func VerifyResources(resources string) error {
	if resources != "" {
		var r map[string]interface{}
		if e := json.Unmarshal([]byte(resources), &r); e != nil {
			return fmt.Errorf("%s: %s should be JSON format", envResource, resources)
		}
	}
	return nil
}

// GetSecret returns the workflow secret name, value.
func GetSecret() (string, string, error) {
	secretMap := make(map[string]map[string]string)
	secretCfg := os.Getenv("SQLFLOW_WORKFLOW_SECRET")
	if secretCfg == "" {
		return "", "", nil
	}
	if e := json.Unmarshal([]byte(secretCfg), &secretMap); e != nil {
		return "", "", e
	}
	if len(secretMap) != 1 {
		return "", "", fmt.Errorf(`SQLFLOW_WORKFLOW_SECRET should be a json string, e.g. {name: {key: value, ...}}`)
	}
	name := reflect.ValueOf(secretMap).MapKeys()[0].String()
	value, e := json.Marshal(secretMap[name])
	if e != nil {
		return "", "", e
	}
	return name, string(value), nil
}

// GenFiller generates Filler to fill the template
func GenFiller(programIR []ir.SQLFlowStmt, session *pb.Session) (*Filler, error) {
	stepEnvs, err := GetStepEnvs(session)
	if err != nil {
		return nil, err
	}
	if os.Getenv("SQLFLOW_WORKFLOW_TTL") != "" {
		workflowTTL, err = strconv.Atoi(os.Getenv("SQLFLOW_WORKFLOW_TTL"))
		if err != nil {
			return nil, fmt.Errorf("SQLFLOW_WORKFLOW_TTL: %s should be int", os.Getenv("SQLFLOW_WORKFLOW_TTL"))
		}
	}
	secretName, secretData, e := GetSecret()
	if e != nil {
		return nil, e
	}
	if e := VerifyResources(os.Getenv(envResource)); e != nil {
		return nil, e
	}

	stepLogFile := os.Getenv("SQLFLOW_WORKFLOW_STEP_LOG_FILE")

	r := &Filler{
		DataSource:  session.DbConnStr,
		StepEnvs:    stepEnvs,
		WorkflowTTL: workflowTTL,
		SecretName:  secretName,
		SecretData:  secretData,
		Resources:   os.Getenv(envResource),
		StepLogFile: stepLogFile,
	}
	// NOTE(yancey1989): does not use ModelImage here since the Predict statement
	// does not contain the ModelImage field in SQL Program IR.
	if os.Getenv("SQLFLOW_WORKFLOW_STEP_IMAGE") != "" {
		defaultDockerImage = os.Getenv("SQLFLOW_WORKFLOW_STEP_IMAGE")
	}

	for _, sqlIR := range programIR {
		switch i := sqlIR.(type) {
		case *ir.NormalStmt, *ir.PredictStmt, *ir.ExplainStmt, *ir.EvaluateStmt:
			// TODO(typhoonzero): get model image used when training.
			sqlStmt := &sqlStatement{
				OriginalSQL: sqlIR.GetOriginalSQL(), IsExtendedSQL: sqlIR.IsExtended(),
				DockerImage: defaultDockerImage}
			r.SQLStatements = append(r.SQLStatements, sqlStmt)
		case *ir.TrainStmt:
			stepImage := defaultDockerImage
			if i.ModelImage != "" {
				stepImage = i.ModelImage
			}
			if r.SQLFlowSubmitter == "katib" {
				sqlStmt, err := ParseKatibSQL(sqlIR.(*ir.TrainStmt))
				if err != nil {
					return nil, fmt.Errorf("Fail to parse Katib train statement %s", sqlIR.GetOriginalSQL())
				}
				r.SQLStatements = append(r.SQLStatements, sqlStmt)
			} else {
				sqlStmt := &sqlStatement{
					OriginalSQL: sqlIR.GetOriginalSQL(), IsExtendedSQL: sqlIR.IsExtended(),
					DockerImage: stepImage}
				r.SQLStatements = append(r.SQLStatements, sqlStmt)
			}
		case *ir.OptimizeStmt:
			sqlStmt := &sqlStatement{
				OriginalSQL:   sqlIR.GetOriginalSQL(),
				IsExtendedSQL: sqlIR.IsExtended(),
				DockerImage:   defaultDockerImage,
			}
			r.SQLStatements = append(r.SQLStatements, sqlStmt)
		case *ir.RunStmt:
			sqlStmt := &sqlStatement{
				OriginalSQL:   sqlIR.GetOriginalSQL(),
				IsExtendedSQL: sqlIR.IsExtended(),
				DockerImage:   i.ImageName,
				Select:        i.Select,
				Into:          i.Into}
			r.SQLStatements = append(r.SQLStatements, sqlStmt)
		default:
			return nil, fmt.Errorf("unrecognized IR type: %v", i)
		}
	}
	return r, nil
}

// GenCode generates a Couler program
func (cg *Codegen) GenCode(programIR []ir.SQLFlowStmt, session *pb.Session) (string, error) {
	r, e := GenFiller(programIR, session)
	if e != nil {
		return "", e
	}
	var program bytes.Buffer
	if err := coulerTemplate.Execute(&program, r); err != nil {
		return "", err
	}
	return program.String(), nil
}

// GenYAML translate the Couler program into Argo YAML
func (cg *Codegen) GenYAML(coulerProgram string) (string, error) {
	cmdline := bytes.Buffer{}
	fmt.Fprintf(&cmdline, "couler run --mode argo --workflow_name sqlflow ")
	if c := os.Getenv("SQLFLOW_WORKFLOW_CLUSTER_CONFIG"); len(c) > 0 {
		fmt.Fprintf(&cmdline, "--cluster_config %s ", c)
	}
	fmt.Fprintf(&cmdline, "--file -")

	coulerExec := strings.Split(cmdline.String(), " ")
	// execute command: `cat sqlflow.couler | couler run --mode argo --workflow_name sqlflow --file -`
	cmd := exec.Command(coulerExec[0], coulerExec[1:]...)
	cmd.Env = append(os.Environ())
	cmd.Stdin = strings.NewReader(coulerProgram)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return "", fmt.Errorf("failed %s, %v %s", cmd, err, out)
	}
	return string(out), nil
}

// MockSQLProgramIR mock a SQLFLow program which contains multiple statements
func MockSQLProgramIR() []ir.SQLFlowStmt {
	normalStmt := ir.NormalStmt("SELECT * FROM iris.train limit 10;")
	trainStmt := ir.MockTrainStmt(true)
	return []ir.SQLFlowStmt{&normalStmt, trainStmt}
}
<commit_msg>Collect step log: add wait to exit time env variable for step (#2905)<commit_after>// Copyright 2020 The SQLFlow Authors. All rights reserved.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package couler

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"reflect"
	"strconv"
	"strings"

	"sqlflow.org/sqlflow/go/ir"
	pb "sqlflow.org/sqlflow/go/proto"
)

var defaultDockerImage = "sqlflow/sqlflow"
var workflowTTL = 24 * 3600
var envResource = "SQLFLOW_WORKFLOW_RESOURCES"

// Codegen generates Couler program
type Codegen struct{}

func fillMapIfValueNotEmpty(m map[string]string, key, value string) {
	if value != "" {
		m[key] = value
	}
}

func fillEnvFromSession(envs *map[string]string, session *pb.Session) {
	fillMapIfValueNotEmpty(*envs, "SQLFLOW_USER_TOKEN", session.Token)
	fillMapIfValueNotEmpty(*envs, "SQLFLOW_DATASOURCE", session.DbConnStr)
	fillMapIfValueNotEmpty(*envs, "SQLFLOW_USER_ID", session.UserId)
	fillMapIfValueNotEmpty(*envs, "SQLFLOW_HIVE_LOCATION", session.HiveLocation)
	fillMapIfValueNotEmpty(*envs, "SQLFLOW_HDFS_NAMENODE_ADDR", session.HdfsNamenodeAddr)
	fillMapIfValueNotEmpty(*envs, "SQLFLOW_HADOOP_USER", session.HdfsUser)
	fillMapIfValueNotEmpty(*envs, "SQLFLOW_HADOOP_PASS", session.HdfsUser)
	fillMapIfValueNotEmpty(*envs, "SQLFLOW_submitter", session.Submitter)
	fillMapIfValueNotEmpty(*envs, "SQLFLOW_WORKFLOW_STEP_LOG_FILE", os.Getenv("SQLFLOW_WORKFLOW_STEP_LOG_FILE"))
	fillMapIfValueNotEmpty(*envs, "SQLFLOW_WORKFLOW_EXIT_TIME_WAIT", os.Getenv("SQLFLOW_WORKFLOW_EXIT_TIME_WAIT"))
}

// GetStepEnvs returns a map of envs used for couler workflow.
func GetStepEnvs(session *pb.Session) (map[string]string, error) {
	envs := make(map[string]string)
	// fill step envs from the environment variables on sqlflowserver
	for _, env := range os.Environ() {
		pair := strings.SplitN(env, "=", 2)
		if len(pair) != 2 {
			return nil, fmt.Errorf("env: %s should format key=value", env)
		}
		// should not pass the workflow env into step
		if strings.HasPrefix(pair[0], "SQLFLOW_") && !strings.HasPrefix(pair[0], "SQLFLOW_WORKFLOW_") {
			envs[pair[0]] = pair[1]
		}
	}
	fillEnvFromSession(&envs, session)
	return envs, nil
}

// VerifyResources verifies the SQLFLOW_WORKFLOW_RESOURCES env to be valid.
func VerifyResources(resources string) error {
	if resources != "" {
		var r map[string]interface{}
		if e := json.Unmarshal([]byte(resources), &r); e != nil {
			return fmt.Errorf("%s: %s should be JSON format", envResource, resources)
		}
	}
	return nil
}

// GetSecret returns the workflow secret name, value.
func GetSecret() (string, string, error) {
	secretMap := make(map[string]map[string]string)
	secretCfg := os.Getenv("SQLFLOW_WORKFLOW_SECRET")
	if secretCfg == "" {
		return "", "", nil
	}
	if e := json.Unmarshal([]byte(secretCfg), &secretMap); e != nil {
		return "", "", e
	}
	if len(secretMap) != 1 {
		return "", "", fmt.Errorf(`SQLFLOW_WORKFLOW_SECRET should be a json string, e.g. {name: {key: value, ...}}`)
	}
	name := reflect.ValueOf(secretMap).MapKeys()[0].String()
	value, e := json.Marshal(secretMap[name])
	if e != nil {
		return "", "", e
	}
	return name, string(value), nil
}

// GenFiller generates Filler to fill the template
func GenFiller(programIR []ir.SQLFlowStmt, session *pb.Session) (*Filler, error) {
	stepEnvs, err := GetStepEnvs(session)
	if err != nil {
		return nil, err
	}
	if os.Getenv("SQLFLOW_WORKFLOW_TTL") != "" {
		workflowTTL, err = strconv.Atoi(os.Getenv("SQLFLOW_WORKFLOW_TTL"))
		if err != nil {
			return nil, fmt.Errorf("SQLFLOW_WORKFLOW_TTL: %s should be int", os.Getenv("SQLFLOW_WORKFLOW_TTL"))
		}
	}
	secretName, secretData, e := GetSecret()
	if e != nil {
		return nil, e
	}
	if e := VerifyResources(os.Getenv(envResource)); e != nil {
		return nil, e
	}

	stepLogFile := os.Getenv("SQLFLOW_WORKFLOW_STEP_LOG_FILE")

	r := &Filler{
		DataSource:  session.DbConnStr,
		StepEnvs:    stepEnvs,
		WorkflowTTL: workflowTTL,
		SecretName:  secretName,
		SecretData:  secretData,
		Resources:   os.Getenv(envResource),
		StepLogFile: stepLogFile,
	}
	// NOTE(yancey1989): does not use ModelImage here since the Predict statement
	// does not contain the ModelImage field in SQL Program IR.
	if os.Getenv("SQLFLOW_WORKFLOW_STEP_IMAGE") != "" {
		defaultDockerImage = os.Getenv("SQLFLOW_WORKFLOW_STEP_IMAGE")
	}

	for _, sqlIR := range programIR {
		switch i := sqlIR.(type) {
		case *ir.NormalStmt, *ir.PredictStmt, *ir.ExplainStmt, *ir.EvaluateStmt:
			// TODO(typhoonzero): get model image used when training.
			sqlStmt := &sqlStatement{
				OriginalSQL: sqlIR.GetOriginalSQL(), IsExtendedSQL: sqlIR.IsExtended(),
				DockerImage: defaultDockerImage}
			r.SQLStatements = append(r.SQLStatements, sqlStmt)
		case *ir.TrainStmt:
			stepImage := defaultDockerImage
			if i.ModelImage != "" {
				stepImage = i.ModelImage
			}
			if r.SQLFlowSubmitter == "katib" {
				sqlStmt, err := ParseKatibSQL(sqlIR.(*ir.TrainStmt))
				if err != nil {
					return nil, fmt.Errorf("Fail to parse Katib train statement %s", sqlIR.GetOriginalSQL())
				}
				r.SQLStatements = append(r.SQLStatements, sqlStmt)
			} else {
				sqlStmt := &sqlStatement{
					OriginalSQL: sqlIR.GetOriginalSQL(), IsExtendedSQL: sqlIR.IsExtended(),
					DockerImage: stepImage}
				r.SQLStatements = append(r.SQLStatements, sqlStmt)
			}
		case *ir.OptimizeStmt:
			sqlStmt := &sqlStatement{
				OriginalSQL:   sqlIR.GetOriginalSQL(),
				IsExtendedSQL: sqlIR.IsExtended(),
				DockerImage:   defaultDockerImage,
			}
			r.SQLStatements = append(r.SQLStatements, sqlStmt)
		case *ir.RunStmt:
			sqlStmt := &sqlStatement{
				OriginalSQL:   sqlIR.GetOriginalSQL(),
				IsExtendedSQL: sqlIR.IsExtended(),
				DockerImage:   i.ImageName,
				Select:        i.Select,
				Into:          i.Into}
			r.SQLStatements = append(r.SQLStatements, sqlStmt)
		default:
			return nil, fmt.Errorf("unrecognized IR type: %v", i)
		}
	}
	return r, nil
}

// GenCode generates a Couler program
func (cg *Codegen) GenCode(programIR []ir.SQLFlowStmt, session *pb.Session) (string, error) {
	r, e := GenFiller(programIR, session)
	if e != nil {
		return "", e
	}
	var program bytes.Buffer
	if err := coulerTemplate.Execute(&program, r); err != nil {
		return "", err
	}
	return program.String(), nil
}

// GenYAML translate the Couler program into Argo YAML
func (cg *Codegen) GenYAML(coulerProgram string) (string, error) {
	cmdline := bytes.Buffer{}
	fmt.Fprintf(&cmdline, "couler run --mode argo --workflow_name sqlflow ")
	if c := os.Getenv("SQLFLOW_WORKFLOW_CLUSTER_CONFIG"); len(c) > 0 {
		fmt.Fprintf(&cmdline, "--cluster_config %s ", c)
	}
	fmt.Fprintf(&cmdline, "--file -")

	coulerExec := strings.Split(cmdline.String(), " ")
	// execute command: `cat sqlflow.couler | couler run --mode argo --workflow_name sqlflow --file -`
	cmd := exec.Command(coulerExec[0], coulerExec[1:]...)
	cmd.Env = append(os.Environ())
	cmd.Stdin = strings.NewReader(coulerProgram)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return "", fmt.Errorf("failed %s, %v %s", cmd, err, out)
	}
	return string(out), nil
}

// MockSQLProgramIR mock a SQLFLow program which contains multiple statements
func MockSQLProgramIR() []ir.SQLFlowStmt {
	normalStmt := ir.NormalStmt("SELECT * FROM iris.train limit 10;")
	trainStmt := ir.MockTrainStmt(true)
	return []ir.SQLFlowStmt{&normalStmt, trainStmt}
}
<|endoftext|>


<commit_before>package goarmorchecksums

import (
	"bytes"
	"crypto/md5"

	"github.com/pkg/errors"
)

var (
	ErrPayloadEmpty   = errors.New("empty payload")
	ErrSecretKeyEmpty = errors.New("empty secret key")
)

func New(toCheck []byte, secretKey string) ([]byte, error) {
	if secretKey == "" {
		return nil, errors.WithStack(ErrSecretKeyEmpty)
	}

	buf := bytes.NewBuffer(toCheck)

	_, err := buf.WriteString(secretKey)
	if err != nil {
		return nil, errors.WithStack(err)
	}

	x := md5.Sum(buf.Bytes())

	if len(toCheck) == 0 {
		return x[:], errors.WithStack(ErrPayloadEmpty)
	}

	return x[:], nil
}

func NewWithSalt(toCheck []byte, secretKey, checksumSalt string) (
	[]byte, error) {
	if secretKey == "" {
		return nil, errors.WithStack(ErrSecretKeyEmpty)
	}

	if checksumSalt == "" {
		return nil, errors.New("missing salt")
	}

	buf := bytes.NewBuffer(toCheck)

	_, err := buf.WriteString(checksumSalt)
	if err != nil {
		return nil, errors.WithStack(err)
	}

	_, err = buf.WriteString(secretKey)
	if err != nil {
		return nil, errors.WithStack(err)
	}

	x := md5.Sum(buf.Bytes())

	if len(toCheck) == 0 {
		return x[:], errors.WithStack(ErrPayloadEmpty)
	}

	return x[:], nil
}
<commit_msg>minor<commit_after>package goarmorchecksums

import (
	"bytes"
	"crypto/md5"
	"encoding/hex"

	"github.com/pkg/errors"
)

var (
	ErrPayloadEmpty   = errors.New("empty payload")
	ErrSecretKeyEmpty = errors.New("empty secret key")
)

func New(toCheck []byte, secretKey string) (string, error) {
	if secretKey == "" {
		return "", errors.WithStack(ErrSecretKeyEmpty)
	}

	buf := bytes.NewBuffer(toCheck)

	_, err := buf.WriteString(secretKey)
	if err != nil {
		return "", errors.WithStack(err)
	}

	a := md5.Sum(buf.Bytes())

	if len(toCheck) == 0 {
		return hex.EncodeToString(a[:]), errors.WithStack(ErrPayloadEmpty)
	}

	return hex.EncodeToString(a[:]), nil
}

func NewWithSalt(toCheck []byte, secretKey, checksumSalt string) (string, error) {
	if secretKey == "" {
		return "", errors.WithStack(ErrSecretKeyEmpty)
	}

	if checksumSalt == "" {
		return "", errors.New("missing salt")
	}

	buf := bytes.NewBuffer(toCheck)

	_, err := buf.WriteString(checksumSalt)
	if err != nil {
		return "", errors.WithStack(err)
	}

	_, err = buf.WriteString(secretKey)
	if err != nil {
		return "", errors.WithStack(err)
	}

	a := md5.Sum(buf.Bytes())

	if len(toCheck) == 0 {
		return hex.EncodeToString(a[:]), errors.WithStack(ErrPayloadEmpty)
	}

	return hex.EncodeToString(a[:]), nil
}
<|endoftext|>


<commit_before>package api

import (
	"fmt"
	"net/http"
	"os"

	"github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"github.com/oskanberg/retrospectiv.es/server/models"
)

// Route represents an API route
type Route struct {
	Name        string
	Method      string
	Pattern     string
	HandlerFunc http.HandlerFunc
}

// Routes respresents multiple routes
type Routes []Route

// BoardAPI in an interface for the BoardAPI commands
type BoardAPI interface {
	Start(string)
}

type api struct {
	boards models.ItemBoards
}

func (a *api) Start(port string) {
	var routes = Routes{
		Route{
			Name:        "Boards",
			Method:      "GET",
			Pattern:     "/api/boards",
			HandlerFunc: a.GetAllBoardsHandler,
		},
		Route{
			Name:        "GetSpecificBoard",
			Method:      "GET",
			Pattern:     "/api/boards/{board_id}",
			HandlerFunc: a.GetSpecificBoardHandler,
		},
		Route{
			Name:        "GetSpecificBoardItems",
			Method:      "GET",
			Pattern:     "/api/boards/{board_id}/items",
			HandlerFunc: a.GetSpecificBoardItemsHandler,
		},
		Route{
			Name:        "AddItemToSpecificBoard",
			Method:      "POST",
			Pattern:     "/api/boards/{board_id}/items",
			HandlerFunc: a.AddItemToSpecificBoard,
		},
		Route{
			Name:        "AddNewBoard",
			Method:      "POST",
			Pattern:     "/api/boards",
			HandlerFunc: a.AddNewBoard,
		},
		Route{
			Name:        "DeleteSpeificItemFromSpecificBoard",
			Method:      "DELETE",
			Pattern:     "/api/boards/{board_id}/items/{item_id}",
			HandlerFunc: a.DeleteSpeificItemFromSpecificBoard,
		},
	}

	r := mux.NewRouter().StrictSlash(true)
	for _, route := range routes {
		r.
			Methods(route.Method).
			Path(route.Pattern).
			Name(route.Name).
			Handler(route.HandlerFunc)
	}

	methods := []string{"GET", "POST", "PUT", "DELETE"}
	log := handlers.LoggingHandler(os.Stdout, r)
	cors := handlers.CORS(handlers.AllowedMethods(methods))(log)
	fmt.Println(http.ListenAndServe(port, cors))
}

// NewBoardAPI returns a pointer to an implementation of BoardAPI
func NewBoardAPI() BoardAPI {
	return &api{
		boards: models.NewItemBoards(),
	}
}
<commit_msg>Add OPTIONS to allowed headers<commit_after>package api

import (
	"fmt"
	"net/http"
	"os"

	"github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"github.com/oskanberg/retrospectiv.es/server/models"
)

// Route represents an API route
type Route struct {
	Name        string
	Method      string
	Pattern     string
	HandlerFunc http.HandlerFunc
}

// Routes respresents multiple routes
type Routes []Route

// BoardAPI in an interface for the BoardAPI commands
type BoardAPI interface {
	Start(string)
}

type api struct {
	boards models.ItemBoards
}

func (a *api) Start(port string) {
	var routes = Routes{
		Route{
			Name:        "Boards",
			Method:      "GET",
			Pattern:     "/api/boards",
			HandlerFunc: a.GetAllBoardsHandler,
		},
		Route{
			Name:        "GetSpecificBoard",
			Method:      "GET",
			Pattern:     "/api/boards/{board_id}",
			HandlerFunc: a.GetSpecificBoardHandler,
		},
		Route{
			Name:        "GetSpecificBoardItems",
			Method:      "GET",
			Pattern:     "/api/boards/{board_id}/items",
			HandlerFunc: a.GetSpecificBoardItemsHandler,
		},
		Route{
			Name:        "AddItemToSpecificBoard",
			Method:      "POST",
			Pattern:     "/api/boards/{board_id}/items",
			HandlerFunc: a.AddItemToSpecificBoard,
		},
		Route{
			Name:        "AddNewBoard",
			Method:      "POST",
			Pattern:     "/api/boards",
			HandlerFunc: a.AddNewBoard,
		},
		Route{
			Name:        "DeleteSpeificItemFromSpecificBoard",
			Method:      "DELETE",
			Pattern:     "/api/boards/{board_id}/items/{item_id}",
			HandlerFunc: a.DeleteSpeificItemFromSpecificBoard,
		},
	}

	r := mux.NewRouter().StrictSlash(true)
	for _, route := range routes {
		r.
			Methods(route.Method).
			Path(route.Pattern).
			Name(route.Name).
			Handler(route.HandlerFunc)
	}

	methods := []string{"GET", "OPTIONS", "POST", "PUT", "DELETE"}
	log := handlers.LoggingHandler(os.Stdout, r)
	cors := handlers.CORS(handlers.AllowedMethods(methods))(log)
	fmt.Println(http.ListenAndServe(port, cors))
}

// NewBoardAPI returns a pointer to an implementation of BoardAPI
func NewBoardAPI() BoardAPI {
	return &api{
		boards: models.NewItemBoards(),
	}
}
<|endoftext|>


<commit_before>package btrfscmd

import (
	"fmt"
	"os/exec"
	"path"
	"strings"
	"bytes"
)

func IsBtrfsVolume(volume string) error {
	cmd := exec.Command("btrfs", "filesystem", "df", volume)
	return cmd.Run()
}

func MakeDeviceBtrfs(device string) error {
	cmd := exec.Command("mkfs.btrfs", "-f", device)
	return cmd.Run()
}

func SubvolumeCreate (volume string, name string) error {
	cmd := exec.Command("btrfs", "subvolume", "create", path.Join(volume, name))
	return cmd.Run()
}

func SubvolumeList (volume string) (result []string) {
	cmd := exec.Command("btrfs", "subvolume", "list", "-a", volume)
	var out bytes.Buffer
	cmd.Stdout = &out
	err := cmd.Run()

	if err != nil {
		fmt.Errorf("Failed to list subvolumes on root volume: %s",volume)
		return
	} else {
		lines := strings.Split(out.String(), "\n")
		for _, line := range lines {
			if tokens := strings.Split(line, "path"); len(tokens) != 2 {
				fmt.Errorf("Can't parse subvolume on line: %s", line)
			} else {
				result = append(result, strings.TrimSpace(tokens[1]))
			}
		}
	}
	return
}

func SubvolumeSnapshot (volume string, sourcesubvolume string, targetsubvolume string, readonly bool) error {
	var cmd *exec.Cmd
	if readonly {
		cmd = exec.Command("btrfs", "subvolume", "snapshot" ,"-r", path.Join(volume,sourcesubvolume), path.Join(volume,targetsubvolume))
	} else {
		cmd = exec.Command("btrfs", "subvolume", "snapshot", path.Join(volume,sourcesubvolume), path.Join(volume,targetsubvolume))
	}
	return cmd.Run()
}

func SubvolumeDelete (volume string, name string) error {
	cmd := exec.Command("btrfs", "subvolume", "delete", path.Join(volume, name))
	return cmd.Run()
}

func Sync() error {
	cmd := exec.Command("sync")
	return cmd.Run()
}

/* TODO: BTRFS Send and Recv */
<commit_msg>Update btrfscmd to list snapshots for volume<commit_after>package btrfscmd

import (
	"bytes"
	"fmt"
	"os/exec"
	"path"
	"strings"
)

func IsBtrfsVolume(volume string) error {
	cmd := exec.Command("btrfs", "filesystem", "df", volume)
	return cmd.Run()
}

func MakeDeviceBtrfs(device string) error {
	cmd := exec.Command("mkfs.btrfs", "-f", device)
	return cmd.Run()
}

func SubvolumeCreate(volume string, name string) error {
	cmd := exec.Command("btrfs", "subvolume", "create", path.Join(volume, name))
	return cmd.Run()
}

func SubvolumeList(volume string) (result []string) {
	cmd := exec.Command("btrfs", "subvolume", "list", "-a", volume)
	var out bytes.Buffer
	cmd.Stdout = &out
	err := cmd.Run()

	if err != nil {
		fmt.Errorf("Failed to list subvolumes on root volume: %s", volume)
		return
	} else {
		lines := strings.Split(out.String(), "\n")
		for _, line := range lines {
			if tokens := strings.Split(line, "path"); len(tokens) != 2 {
				fmt.Errorf("Can't parse subvolume on line: %s", line)
			} else {
				result = append(result, strings.TrimSpace(tokens[1]))
			}
		}
	}
	return
}

func SubvolumeSnapshot(volume string, sourcesubvolume string, targetsubvolume string, readonly bool) error {
	var cmd *exec.Cmd
	if readonly {
		cmd = exec.Command("btrfs", "subvolume", "snapshot", "-r", path.Join(volume, sourcesubvolume), path.Join(volume, targetsubvolume))
	} else {
		cmd = exec.Command("btrfs", "subvolume", "snapshot", path.Join(volume, sourcesubvolume), path.Join(volume, targetsubvolume))
	}
	return cmd.Run()
}

func SubvolumeDelete(volume string, name string) error {
	cmd := exec.Command("btrfs", "subvolume", "delete", path.Join(volume, name))
	return cmd.Run()
}

func Sync() error {
	cmd := exec.Command("sync")
	return cmd.Run()
}

/* Gets the list of snapshot for a parent subvolume*/
func SubvolumeSnapshotList(volume string, parentsubvolume string) (result []string) {
	cmd := exec.Command("btrfs", "subvolume", "show", path.Join(volume, parentsubvolume))
	var out bytes.Buffer
	cmd.Stdout = &out
	err := cmd.Run()

	if err != nil {
		fmt.Errorf("Failed to get subvolume information: %s", path.Join(volume, parentsubvolume))
		return
	} else {
		lines := strings.Split(out.String(), "\n")
		parent_uuid := strings.TrimSpace(strings.Split(lines[2], "uuid:")[1])
		cmd2 := exec.Command("btrfs", "subvolume", "list", "-qus", path.Join(volume, parentsubvolume))
		var out2 bytes.Buffer
		cmd2.Stdout = &out2
		err2 := cmd2.Run()
		if err2 != nil {
			fmt.Errorf("Failed to list snapshots on subvolume : %s", path.Join(volume, parentsubvolume))
		} else {
			lines2 := strings.Split(out2.String(), "\n")
			for _, line2 := range lines2 {
				tokens := strings.Split(line2, " ")
				if len(tokens)==18 && tokens[13] == parent_uuid {
					result = append(result, tokens[17])
				}
			}
		}
	}
	return
}

/* TODO: BTRFS Send and Recv */
<|endoftext|>


<commit_before>package msg
import (
	"github.com/name5566/leaf/network/protobuf"
)

var (
	Processor = protobuf.NewProcessor()
)

func init() {	// 这里我们注册了一个 protobuf 消息)
    Processor.Register(&StartFight{})
    Processor.Register(&FightResult{})
    Processor.Register(&EnterFight{})
    Processor.Register(&SignUpResponse{})
    Processor.Register(&TosChat{})
    Processor.Register(&TocChat{})
    Processor.Register(&Login{})
    Processor.Register(&PlayerBaseInfo{})
    Processor.Register(&LoginSuccessfull{})
    Processor.Register(&LoginFaild{})

}<commit_msg>Update msg.go<commit_after>package msg
import (
	"github.com/name5566/leaf/network/protobuf"
)

var (
	Processor = protobuf.NewProcessor()
)

func init() {	// 这里我们注册了一个 protobuf 消息)
    Processor.SetByteOrder(true)
    Processor.Register(&StartFight{})
    Processor.Register(&FightResult{})
    Processor.Register(&EnterFight{})
    Processor.Register(&SignUpResponse{})
    Processor.Register(&TosChat{})
    Processor.Register(&TocChat{})
    Processor.Register(&Login{})
    Processor.Register(&PlayerBaseInfo{})
    Processor.Register(&LoginSuccessfull{})
    Processor.Register(&LoginFaild{})

}
<|endoftext|>


<commit_before>/*
   Copyright The containerd Authors.

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

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

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

package shim

import (
	"context"
	"fmt"
	"net"
	"os"
	"syscall"
	"time"

	winio "github.com/Microsoft/go-winio"
	"github.com/containerd/containerd/namespaces"
	"github.com/pkg/errors"
)

func getSysProcAttr() *syscall.SysProcAttr {
	return nil
}

// SetScore sets the oom score for a process
func SetScore(pid int) error {
	return nil
}

// SocketAddress returns a npipe address
func SocketAddress(ctx context.Context, id string) (string, error) {
	ns, err := namespaces.NamespaceRequired(ctx)
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("\\\\.\\pipe\\containerd-shim-%s-%s-pipe", ns, id), nil
}

// AnonDialer returns a dialer for a npipe
func AnonDialer(address string, timeout time.Duration) (net.Conn, error) {
	var c net.Conn
	var lastError error
	start := time.Now()
	for {
		remaining := timeout - time.Now().Sub(start)
		if remaining <= 0 {
			lastError = errors.Errorf("timed out waiting for npipe %s", address)
			break
		}
		c, lastError = winio.DialPipe(address, &remaining)
		if lastError == nil {
			break
		}
		if !os.IsNotExist(lastError) {
			break
		}
		time.Sleep(10 * time.Millisecond)
	}
	return c, lastError
}

// NewSocket returns a new npipe listener
func NewSocket(address string) (net.Listener, error) {
	l, err := winio.ListenPipe(address, nil)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to listen to npipe %s", address)
	}
	return l, nil
}
<commit_msg>Decrease shim timeout on pipe not found<commit_after>/*
   Copyright The containerd Authors.

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

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

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

package shim

import (
	"context"
	"fmt"
	"net"
	"os"
	"syscall"
	"time"

	winio "github.com/Microsoft/go-winio"
	"github.com/containerd/containerd/namespaces"
	"github.com/pkg/errors"
)

func getSysProcAttr() *syscall.SysProcAttr {
	return nil
}

// SetScore sets the oom score for a process
func SetScore(pid int) error {
	return nil
}

// SocketAddress returns a npipe address
func SocketAddress(ctx context.Context, id string) (string, error) {
	ns, err := namespaces.NamespaceRequired(ctx)
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("\\\\.\\pipe\\containerd-shim-%s-%s-pipe", ns, id), nil
}

// AnonDialer returns a dialer for a npipe
func AnonDialer(address string, timeout time.Duration) (net.Conn, error) {
	var c net.Conn
	var lastError error
	timedOutError := errors.Errorf("timed out waiting for npipe %s", address)
	start := time.Now()
	for {
		remaining := timeout - time.Now().Sub(start)
		if remaining <= 0 {
			lastError = timedOutError
			break
		}
		c, lastError = winio.DialPipe(address, &remaining)
		if lastError == nil {
			break
		}
		if !os.IsNotExist(lastError) {
			break
		}
		// There is nobody serving the pipe. We limit the timeout for this case
		// to 5 seconds because any shim that would serve this endpoint should
		// serve it within 5 seconds. We use the passed in timeout for the
		// `DialPipe` timeout if the pipe exists however to give the pipe time
		// to `Accept` the connection.
		if time.Now().Sub(start) >= 5*time.Second {
			lastError = timedOutError
			break
		}
		time.Sleep(10 * time.Millisecond)
	}
	return c, lastError
}

// NewSocket returns a new npipe listener
func NewSocket(address string) (net.Listener, error) {
	l, err := winio.ListenPipe(address, nil)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to listen to npipe %s", address)
	}
	return l, nil
}
<|endoftext|>


<commit_before><commit_msg>Fix rebase removed pointer for SupportUpgrade<commit_after><|endoftext|>


<commit_before>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package phases

import (
	"errors"
	"os"

	"k8s.io/kubernetes/cmd/kubeadm/app/cmd/phases/workflow"
	kubeadmconstants "k8s.io/kubernetes/cmd/kubeadm/app/constants"
	"k8s.io/kubernetes/cmd/kubeadm/app/phases/uploadconfig"
)

// NewUpdateClusterStatus creates a kubeadm workflow phase for update-cluster-status
func NewUpdateClusterStatus() workflow.Phase {
	return workflow.Phase{
		Name:  "update-cluster-status",
		Short: "Remove this node from the ClusterStatus object.",
		Long:  "Remove this node from the ClusterStatus object if the node is a control plane node.",
		Run:   runUpdateClusterStatus,
	}
}

func runUpdateClusterStatus(c workflow.RunData) error {
	r, ok := c.(resetData)
	if !ok {
		return errors.New("update-cluster-status phase invoked with an invalid data struct")
	}

	// Reset the ClusterStatus for a given control-plane node.
	cfg := r.Cfg()
	if isControlPlane() && cfg != nil {
		uploadconfig.ResetClusterStatusForNode(cfg.NodeRegistration.Name, r.Client())
	}

	return nil
}

// isControlPlane checks if a node is a control-plane node by looking up
// the kube-apiserver manifest file
func isControlPlane() bool {
	filepath := kubeadmconstants.GetStaticPodFilepath(kubeadmconstants.KubeAPIServer, kubeadmconstants.GetStaticPodDirectory())
	_, err := os.Stat(filepath)
	return !os.IsNotExist(err)
}
<commit_msg>kubeadm: handle ResetClusterStatusForNode errors<commit_after>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package phases

import (
	"errors"
	"os"

	"k8s.io/kubernetes/cmd/kubeadm/app/cmd/phases/workflow"
	kubeadmconstants "k8s.io/kubernetes/cmd/kubeadm/app/constants"
	"k8s.io/kubernetes/cmd/kubeadm/app/phases/uploadconfig"
)

// NewUpdateClusterStatus creates a kubeadm workflow phase for update-cluster-status
func NewUpdateClusterStatus() workflow.Phase {
	return workflow.Phase{
		Name:  "update-cluster-status",
		Short: "Remove this node from the ClusterStatus object.",
		Long:  "Remove this node from the ClusterStatus object if the node is a control plane node.",
		Run:   runUpdateClusterStatus,
	}
}

func runUpdateClusterStatus(c workflow.RunData) error {
	r, ok := c.(resetData)
	if !ok {
		return errors.New("update-cluster-status phase invoked with an invalid data struct")
	}

	// Reset the ClusterStatus for a given control-plane node.
	cfg := r.Cfg()
	if isControlPlane() && cfg != nil {
		if err := uploadconfig.ResetClusterStatusForNode(cfg.NodeRegistration.Name, r.Client()); err != nil {
			return err
		}
	}

	return nil
}

// isControlPlane checks if a node is a control-plane node by looking up
// the kube-apiserver manifest file
func isControlPlane() bool {
	filepath := kubeadmconstants.GetStaticPodFilepath(kubeadmconstants.KubeAPIServer, kubeadmconstants.GetStaticPodDirectory())
	_, err := os.Stat(filepath)
	return !os.IsNotExist(err)
}
<|endoftext|>


<commit_before>// archive is package that helps create archives in a format that
// Harmony expects with its various upload endpoints.
package archive

import (
	"archive/tar"
	"bufio"
	"compress/gzip"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
)

// ArchiveOpts are the options for defining how the archive will be built.
type ArchiveOpts struct {
	// Exclude and Include are filters of files to include/exclude in
	// the archive when creating it from a directory. These filters should
	// be relative to the packaging directory and should be basic glob
	// patterns.
	Exclude []string
	Include []string

	// Extra is a mapping of extra files to include within the archive. The
	// key should be the path within the archive and the value should be
	// an absolute path to the file to put into the archive. These extra
	// files will override any other files in the archive.
	Extra map[string]string

	// VCS, if true, will detect and use a VCS system to determien what
	// files to include the archive.
	VCS bool
}

// IsSet says whether any options were set.
func (o *ArchiveOpts) IsSet() bool {
	return len(o.Exclude) > 0 || len(o.Include) > 0 || o.VCS
}

// Archive takes the given path and ArchiveOpts and archives it.
//
// The archive is done async and streamed via the io.ReadCloser returned.
// The reader is blocking: data is only compressed and written as data is
// being read from the reader. Because of this, any user doesn't have to
// worry about quickly reading data to avoid memory bloat.
//
// The archive can be read with the io.ReadCloser that is returned. The error
// returned is an error that happened before archiving started, so the
// ReadCloser doesn't need to be closed (and should be nil). The error
// channel are errors that can happen while archiving is happening. When
// an error occurs on the channel, reading should stop and be closed.
func Archive(
	path string, opts *ArchiveOpts) (io.ReadCloser, <-chan error, error) {
	fi, err := os.Stat(path)
	if err != nil {
		return nil, nil, err
	}

	// Direct file paths cannot have archive options
	if !fi.IsDir() && opts.IsSet() {
		return nil, nil, fmt.Errorf(
			"Options such as exclude, include, and VCS can't be set when " +
				"the path is a file.")
	}

	if fi.IsDir() {
		return archiveDir(path, opts)
	} else {
		return archiveFile(path, opts)
	}
}

func archiveFile(
	path string, opts *ArchiveOpts) (io.ReadCloser, <-chan error, error) {
	// TODO: if file is already gzipped, then send it along
	// TODO: if file is not gzipped, then... error? or do we tar + gzip?

	return nil, nil, nil
}

func archiveDir(
	root string, opts *ArchiveOpts) (io.ReadCloser, <-chan error, error) {
	var vcsInclude []string
	if opts.VCS {
		var err error
		vcsInclude, err = vcsFiles(root)
		if err != nil {
			return nil, nil, err
		}
	}

	// We're going to write to an io.Pipe so that we can ensure the other
	// side is reading as we're writing.
	pr, pw := io.Pipe()

	// Buffer the writer so that we can keep some data moving in memory
	// while we're compressing. 4M should be good.
	bufW := bufio.NewWriterSize(pw, 4096*1024)

	// Gzip compress all the output data
	gzipW := gzip.NewWriter(bufW)

	// Tar the file contents
	tarW := tar.NewWriter(gzipW)

	// Build the function that'll do all the compression
	walkFn := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// Get the relative path from the path since it contains the root
		// plus the path.
		subpath, err := filepath.Rel(root, path)
		if err != nil {
			return err
		}
		if subpath == "." {
			return nil
		}

		// If we have a list of VCS files, check that first
		skip := false
		if len(vcsInclude) > 0 {
			skip = true
			for _, f := range vcsInclude {
				if f == subpath {
					skip = false
					break
				}

				if info.IsDir() && strings.HasPrefix(f, subpath+"/") {
					skip = false
					break
				}
			}
		}

		// If include is present, we only include what is listed
		if len(opts.Include) > 0 {
			skip = true
			for _, include := range opts.Include {
				match, err := filepath.Match(include, subpath)
				if err != nil {
					return err
				}
				if match {
					skip = false
					break
				}
			}
		}

		// If exclude, it is one last gate to excluding files
		for _, exclude := range opts.Exclude {
			match, err := filepath.Match(exclude, subpath)
			if err != nil {
				return err
			}
			if match {
				skip = true
				break
			}
		}

		// If we have to skip this file, then skip it, properly skipping
		// children if we're a directory.
		if skip {
			if info.IsDir() {
				return filepath.SkipDir
			}

			return nil
		}

		// Read the symlink target. We don't track the error because
		// it doesn't matter if there is an error.
		target, _ := os.Readlink(path)

		// Build the file header for the tar entry
		header, err := tar.FileInfoHeader(info, target)
		if err != nil {
			return fmt.Errorf(
				"Failed creating archive header: %s", path)
		}

		// Modify the header to properly be the full subpath
		header.Name = subpath
		if info.IsDir() {
			header.Name += "/"
		}

		// Write the header first to the archive.
		if err := tarW.WriteHeader(header); err != nil {
			return fmt.Errorf(
				"Failed writing archive header: %s", path)
		}

		// If it is a directory, then we're done (no body to write)
		if info.IsDir() {
			return nil
		}

		// Open the target file to write the data
		f, err := os.Open(path)
		if err != nil {
			return fmt.Errorf(
				"Failed opening file '%s' to write compressed archive.", path)
		}
		defer f.Close()

		if _, err = io.Copy(tarW, f); err != nil {
			return fmt.Errorf(
				"Failed copying file to archive: %s", path)
		}

		return nil
	}

	// Create all our channels so we can send data through some tubes
	// to other goroutines.
	errCh := make(chan error, 1)
	go func() {
		// First, walk the path and do the normal files
		werr := filepath.Walk(root, walkFn)
		if werr != nil {
			goto CLEANUP
		}

		// If that succeeded, handle the extra files
		werr = copyExtras(tarW, opts.Extra)

	CLEANUP:
		// Attempt to close all the things. If we get an error on the way
		// and we haven't had an error yet, then record that as the critical
		// error. But we still try to close everything.

		// Close the tar writer
		if err := tarW.Close(); err != nil && werr == nil {
			werr = err
		}

		// Close the gzip writer
		if err := gzipW.Close(); err != nil && werr == nil {
			werr = err
		}

		// Flush the buffer
		if err := bufW.Flush(); err != nil && werr == nil {
			werr = err
		}

		// Close the pipe
		if err := pw.Close(); err != nil && werr == nil {
			werr = err
		}

		// Send any error we might have down the pipe if we have one
		if werr != nil {
			errCh <- werr
		}
	}()

	return pr, errCh, nil
}

func copyExtras(w *tar.Writer, extra map[string]string) error {
	for entry, path := range extra {
		info, err := os.Stat(path)
		if err != nil {
			return err
		}

		// Read the symlink target. We don't track the error because
		// it doesn't matter if there is an error.
		target, _ := os.Readlink(path)

		// Build the file header for the tar entry
		header, err := tar.FileInfoHeader(info, target)
		if err != nil {
			return fmt.Errorf(
				"Failed creating archive header: %s", path)
		}

		// Modify the header to properly be the full subpath
		header.Name = entry

		// Write the header first to the archive.
		if err := w.WriteHeader(header); err != nil {
			return fmt.Errorf(
				"Failed writing archive header: %s", path)
		}

		// If it is a directory, then we're done (no body to write)
		if info.IsDir() {
			return nil
		}

		// Open the target file to write the data
		f, err := os.Open(path)
		if err != nil {
			return fmt.Errorf(
				"Failed opening file '%s' to write compressed archive.", path)
		}
		defer f.Close()

		if _, err = io.Copy(w, f); err != nil {
			return fmt.Errorf(
				"Failed copying file to archive: %s", path)
		}
	}

	return nil
}
<commit_msg>archive: close files right away for extras<commit_after>// archive is package that helps create archives in a format that
// Harmony expects with its various upload endpoints.
package archive

import (
	"archive/tar"
	"bufio"
	"compress/gzip"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
)

// ArchiveOpts are the options for defining how the archive will be built.
type ArchiveOpts struct {
	// Exclude and Include are filters of files to include/exclude in
	// the archive when creating it from a directory. These filters should
	// be relative to the packaging directory and should be basic glob
	// patterns.
	Exclude []string
	Include []string

	// Extra is a mapping of extra files to include within the archive. The
	// key should be the path within the archive and the value should be
	// an absolute path to the file to put into the archive. These extra
	// files will override any other files in the archive.
	Extra map[string]string

	// VCS, if true, will detect and use a VCS system to determien what
	// files to include the archive.
	VCS bool
}

// IsSet says whether any options were set.
func (o *ArchiveOpts) IsSet() bool {
	return len(o.Exclude) > 0 || len(o.Include) > 0 || o.VCS
}

// Archive takes the given path and ArchiveOpts and archives it.
//
// The archive is done async and streamed via the io.ReadCloser returned.
// The reader is blocking: data is only compressed and written as data is
// being read from the reader. Because of this, any user doesn't have to
// worry about quickly reading data to avoid memory bloat.
//
// The archive can be read with the io.ReadCloser that is returned. The error
// returned is an error that happened before archiving started, so the
// ReadCloser doesn't need to be closed (and should be nil). The error
// channel are errors that can happen while archiving is happening. When
// an error occurs on the channel, reading should stop and be closed.
func Archive(
	path string, opts *ArchiveOpts) (io.ReadCloser, <-chan error, error) {
	fi, err := os.Stat(path)
	if err != nil {
		return nil, nil, err
	}

	// Direct file paths cannot have archive options
	if !fi.IsDir() && opts.IsSet() {
		return nil, nil, fmt.Errorf(
			"Options such as exclude, include, and VCS can't be set when " +
				"the path is a file.")
	}

	if fi.IsDir() {
		return archiveDir(path, opts)
	} else {
		return archiveFile(path, opts)
	}
}

func archiveFile(
	path string, opts *ArchiveOpts) (io.ReadCloser, <-chan error, error) {
	// TODO: if file is already gzipped, then send it along
	// TODO: if file is not gzipped, then... error? or do we tar + gzip?

	return nil, nil, nil
}

func archiveDir(
	root string, opts *ArchiveOpts) (io.ReadCloser, <-chan error, error) {
	var vcsInclude []string
	if opts.VCS {
		var err error
		vcsInclude, err = vcsFiles(root)
		if err != nil {
			return nil, nil, err
		}
	}

	// We're going to write to an io.Pipe so that we can ensure the other
	// side is reading as we're writing.
	pr, pw := io.Pipe()

	// Buffer the writer so that we can keep some data moving in memory
	// while we're compressing. 4M should be good.
	bufW := bufio.NewWriterSize(pw, 4096*1024)

	// Gzip compress all the output data
	gzipW := gzip.NewWriter(bufW)

	// Tar the file contents
	tarW := tar.NewWriter(gzipW)

	// Build the function that'll do all the compression
	walkFn := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// Get the relative path from the path since it contains the root
		// plus the path.
		subpath, err := filepath.Rel(root, path)
		if err != nil {
			return err
		}
		if subpath == "." {
			return nil
		}

		// If we have a list of VCS files, check that first
		skip := false
		if len(vcsInclude) > 0 {
			skip = true
			for _, f := range vcsInclude {
				if f == subpath {
					skip = false
					break
				}

				if info.IsDir() && strings.HasPrefix(f, subpath+"/") {
					skip = false
					break
				}
			}
		}

		// If include is present, we only include what is listed
		if len(opts.Include) > 0 {
			skip = true
			for _, include := range opts.Include {
				match, err := filepath.Match(include, subpath)
				if err != nil {
					return err
				}
				if match {
					skip = false
					break
				}
			}
		}

		// If exclude, it is one last gate to excluding files
		for _, exclude := range opts.Exclude {
			match, err := filepath.Match(exclude, subpath)
			if err != nil {
				return err
			}
			if match {
				skip = true
				break
			}
		}

		// If we have to skip this file, then skip it, properly skipping
		// children if we're a directory.
		if skip {
			if info.IsDir() {
				return filepath.SkipDir
			}

			return nil
		}

		// Read the symlink target. We don't track the error because
		// it doesn't matter if there is an error.
		target, _ := os.Readlink(path)

		// Build the file header for the tar entry
		header, err := tar.FileInfoHeader(info, target)
		if err != nil {
			return fmt.Errorf(
				"Failed creating archive header: %s", path)
		}

		// Modify the header to properly be the full subpath
		header.Name = subpath
		if info.IsDir() {
			header.Name += "/"
		}

		// Write the header first to the archive.
		if err := tarW.WriteHeader(header); err != nil {
			return fmt.Errorf(
				"Failed writing archive header: %s", path)
		}

		// If it is a directory, then we're done (no body to write)
		if info.IsDir() {
			return nil
		}

		// Open the target file to write the data
		f, err := os.Open(path)
		if err != nil {
			return fmt.Errorf(
				"Failed opening file '%s' to write compressed archive.", path)
		}
		defer f.Close()

		if _, err = io.Copy(tarW, f); err != nil {
			return fmt.Errorf(
				"Failed copying file to archive: %s", path)
		}

		return nil
	}

	// Create all our channels so we can send data through some tubes
	// to other goroutines.
	errCh := make(chan error, 1)
	go func() {
		// First, walk the path and do the normal files
		werr := filepath.Walk(root, walkFn)
		if werr != nil {
			goto CLEANUP
		}

		// If that succeeded, handle the extra files
		werr = copyExtras(tarW, opts.Extra)

	CLEANUP:
		// Attempt to close all the things. If we get an error on the way
		// and we haven't had an error yet, then record that as the critical
		// error. But we still try to close everything.

		// Close the tar writer
		if err := tarW.Close(); err != nil && werr == nil {
			werr = err
		}

		// Close the gzip writer
		if err := gzipW.Close(); err != nil && werr == nil {
			werr = err
		}

		// Flush the buffer
		if err := bufW.Flush(); err != nil && werr == nil {
			werr = err
		}

		// Close the pipe
		if err := pw.Close(); err != nil && werr == nil {
			werr = err
		}

		// Send any error we might have down the pipe if we have one
		if werr != nil {
			errCh <- werr
		}
	}()

	return pr, errCh, nil
}

func copyExtras(w *tar.Writer, extra map[string]string) error {
	for entry, path := range extra {
		info, err := os.Stat(path)
		if err != nil {
			return err
		}

		// Read the symlink target. We don't track the error because
		// it doesn't matter if there is an error.
		target, _ := os.Readlink(path)

		// Build the file header for the tar entry
		header, err := tar.FileInfoHeader(info, target)
		if err != nil {
			return fmt.Errorf(
				"Failed creating archive header: %s", path)
		}

		// Modify the header to properly be the full subpath
		header.Name = entry

		// Write the header first to the archive.
		if err := w.WriteHeader(header); err != nil {
			return fmt.Errorf(
				"Failed writing archive header: %s", path)
		}

		// If it is a directory, then we're done (no body to write)
		if info.IsDir() {
			return nil
		}

		// Open the target file to write the data
		f, err := os.Open(path)
		if err != nil {
			return fmt.Errorf(
				"Failed opening file '%s' to write compressed archive.", path)
		}

		_, err = io.Copy(w, f)
		f.Close()
		if err != nil {
			return fmt.Errorf(
				"Failed copying file to archive: %s", path)
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package fp_reg

// host[:port] of the facilitator you want to register with
// for example, fp-facilitator.org
const FP_FACILITATOR = ""
<commit_msg>FP_FACILITATOR can contain a base path too<commit_after>package fp_reg

// host:port/basepath of the facilitator you want to register with
// for example, fp-facilitator.org or example.com:12345/facilitator
// https:// and /reg/ will be prepended and appended respectively.
const FP_FACILITATOR = ""
<|endoftext|>