<commit_before>/*
Copyright 2021 The Vitess 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 grpctmclient

import (
	"context"
	"flag"
	"io"
	"sort"
	"sync"
	"time"

	"google.golang.org/grpc"

	"vitess.io/vitess/go/netutil"
	"vitess.io/vitess/go/stats"
	"vitess.io/vitess/go/sync2"
	"vitess.io/vitess/go/vt/grpcclient"
	"vitess.io/vitess/go/vt/vttablet/tmclient"

	tabletmanagerservicepb "vitess.io/vitess/go/vt/proto/tabletmanagerservice"
	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
)

var (
	defaultPoolCapacity = flag.Int("tablet_manager_grpc_connpool_size", 100, "number of tablets to keep tmclient connections open to")
)

func init() {
	tmclient.RegisterTabletManagerClientFactory("grpc-cached", func() tmclient.TabletManagerClient {
		return NewCachedConnClient(*defaultPoolCapacity)
	})
}

// closeFunc allows a standalone function to implement io.Closer, similar to
// how http.HandlerFunc allows standalone functions to implement http.Handler.
type closeFunc func() error

func (fn closeFunc) Close() error {
	return fn()
}

var _ io.Closer = (*closeFunc)(nil)

type cachedConn struct {
	tabletmanagerservicepb.TabletManagerClient
	cc *grpc.ClientConn

	addr           string
	lastAccessTime time.Time
	refs           int
}

type cachedConnDialer struct {
	m            sync.Mutex
	conns        map[string]*cachedConn
	evict        []*cachedConn
	evictSorted  bool
	connWaitSema *sync2.Semaphore
}

var dialerStats = struct {
	ConnReuse    *stats.Gauge
	ConnNew      *stats.Gauge
	DialTimeouts *stats.Gauge
	DialTimings  *stats.Timings
}{
	ConnReuse:    stats.NewGauge("tabletmanagerclient_cachedconn_reuse", "number of times a call to dial() was able to reuse an existing connection"),
	ConnNew:      stats.NewGauge("tabletmanagerclient_cachedconn_new", "number of times a call to dial() resulted in a dialing a new grpc clientconn"),
	DialTimeouts: stats.NewGauge("tabletmanagerclient_cachedconn_dial_timeouts", "number of context timeouts during dial()"),
	DialTimings:  stats.NewTimings("tabletmanagerclient_cachedconn_dialtimings", "timings for various dial paths", "path", "rlock_fast", "sema_fast", "sema_poll"),
}

// NewCachedConnClient returns a grpc Client that caches connections to the different tablets
func NewCachedConnClient(capacity int) *Client {
	dialer := &cachedConnDialer{
		conns:        make(map[string]*cachedConn, capacity),
		evict:        make([]*cachedConn, 0, capacity),
		connWaitSema: sync2.NewSemaphore(capacity, 0),
	}
	return &Client{dialer}
}

var _ dialer = (*cachedConnDialer)(nil)

func (dialer *cachedConnDialer) sortEvictionsLocked() {
	if !dialer.evictSorted {
		sort.Slice(dialer.evict, func(i, j int) bool {
			left, right := dialer.evict[i], dialer.evict[j]
			if left.refs == right.refs {
				return left.lastAccessTime.After(right.lastAccessTime)
			}
			return left.refs > right.refs
		})
		dialer.evictSorted = true
	}
}

func (dialer *cachedConnDialer) dial(ctx context.Context, tablet *topodatapb.Tablet) (tabletmanagerservicepb.TabletManagerClient, io.Closer, error) {
	start := time.Now()
	addr := getTabletAddr(tablet)

	if client, closer, found, err := dialer.tryFromCache(addr, &dialer.m); found {
		dialerStats.DialTimings.Add("rlock_fast", time.Since(start))
		return client, closer, err
	}

	if dialer.connWaitSema.TryAcquire() {
		defer func() {
			dialerStats.DialTimings.Add("sema_fast", time.Since(start))
		}()

		// Check if another goroutine managed to dial a conn for the same addr
		// while we were waiting for the write lock. This is identical to the
		// read-lock section above.
		if client, closer, found, err := dialer.tryFromCache(addr, &dialer.m); found {
			return client, closer, err
		}
		return dialer.newdial(addr)
	}

	defer func() {
		dialerStats.DialTimings.Add("sema_poll", time.Since(start))
	}()

	for {
		select {
		case <-ctx.Done():
			dialerStats.DialTimeouts.Add(1)
			return nil, nil, ctx.Err()
		default:
			if client, closer, found, err := dialer.pollOnce(addr); found {
				return client, closer, err
			}
		}
	}
}

// tryFromCache tries to get a connection from the cache, performing a redial
// on that connection if it exists. It returns a TabletManagerClient impl, an
// io.Closer, a flag to indicate whether a connection was found in the cache,
// and an error, which is always nil.
//
// In addition to the addr being dialed, tryFromCache takes a sync.Locker which,
// if not nil, will be used to wrap the lookup and redial in that lock. This
// function can be called in situations where the conns map is locked
// externally (like in pollOnce), so we do not want to manage the locks here. In
// other cases (like in the rlock_fast path of dial()), we pass in the RLocker
// to ensure we have a read lock on the cache for the duration of the call.
func (dialer *cachedConnDialer) tryFromCache(addr string, locker sync.Locker) (client tabletmanagerservicepb.TabletManagerClient, closer io.Closer, found bool, err error) {
	if locker != nil {
		locker.Lock()
		defer locker.Unlock()
	}

	if conn, ok := dialer.conns[addr]; ok {
		client, closer, err := dialer.redialLocked(conn)
		return client, closer, ok, err
	}

	return nil, nil, false, nil
}

// pollOnce is called on each iteration of the polling loop in dial(). It:
// - locks the conns cache for writes
// - attempts to get a connection from the cache. If found, redial() it and exit.
// - locks the queue
// - peeks at the head of the eviction queue. if the peeked conn has no refs, it
//   is unused, and can be evicted to make room for the new connection to addr.
//   If the peeked conn has refs, exit.
// - pops the conn we just peeked from the queue, delete it from the cache, and
//   close the underlying ClientConn for that conn.
// - attempt a newdial. if the newdial fails, it will release a slot on the
//   connWaitSema, so another dial() call can successfully acquire it to dial
//   a new conn. if the newdial succeeds, we will have evicted one conn, but
//   added another, so the net change is 0, and no changes to the connWaitSema
//   are made.
//
// It returns a TabletManagerClient impl, an io.Closer, a flag to indicate
// whether the dial() poll loop should exit, and an error.
func (dialer *cachedConnDialer) pollOnce(addr string) (client tabletmanagerservicepb.TabletManagerClient, closer io.Closer, found bool, err error) {
	dialer.m.Lock()

	if client, closer, found, err := dialer.tryFromCache(addr, nil); found {
		dialer.m.Unlock()
		return client, closer, found, err
	}

	dialer.sortEvictionsLocked()

	conn := dialer.evict[len(dialer.evict)-1]
	if conn.refs != 0 {
		dialer.m.Unlock()
		return nil, nil, false, nil
	}

	dialer.evict = dialer.evict[:len(dialer.evict)-1]
	delete(dialer.conns, conn.addr)
	conn.cc.Close()
	dialer.m.Unlock()

	client, closer, err = dialer.newdial(addr)
	return client, closer, true, err
}

// newdial creates a new cached connection, and updates the cache and eviction
// queue accordingly. If newdial fails to create the underlying
// gRPC connection, it will make a call to Release the connWaitSema for other
// newdial calls.
//
// It returns the three-tuple of client-interface, closer, and error that the
// main dial func returns.
func (dialer *cachedConnDialer) newdial(addr string) (tabletmanagerservicepb.TabletManagerClient, io.Closer, error) {
	dialerStats.ConnNew.Add(1)

	opt, err := grpcclient.SecureDialOption(*cert, *key, *ca, *name)
	if err != nil {
		dialer.connWaitSema.Release()
		return nil, nil, err
	}

	cc, err := grpcclient.Dial(addr, grpcclient.FailFast(false), opt)
	if err != nil {
		dialer.connWaitSema.Release()
		return nil, nil, err
	}

	dialer.m.Lock()
	defer dialer.m.Unlock()

	if conn, existing := dialer.conns[addr]; existing {
		// race condition: some other goroutine has dialed our tablet before we have;
		// this is not great, but shouldn't happen often (if at all), so we're going to
		// close this connection and reuse the existing one. by doing this, we can keep
		// the actual Dial out of the global lock and significantly increase throughput
		cc.Close()
		return dialer.redialLocked(conn)
	}

	conn := &cachedConn{
		TabletManagerClient: tabletmanagerservicepb.NewTabletManagerClient(cc),
		cc:                  cc,
		lastAccessTime:      time.Now(),
		refs:                1,
		addr:                addr,
	}
	dialer.evict = append(dialer.evict, conn)
	dialer.evictSorted = false
	dialer.conns[addr] = conn

	return dialer.connWithCloser(conn)
}

// redialLocked takes an already-dialed connection in the cache does all the work of
// lending that connection out to one more caller. this should only ever be
// called while holding at least the RLock on dialer.m (but the write lock is
// fine too), to prevent the connection from getting evicted out from under us.
//
// It returns the three-tuple of client-interface, closer, and error that the
// main dial func returns.
func (dialer *cachedConnDialer) redialLocked(conn *cachedConn) (tabletmanagerservicepb.TabletManagerClient, io.Closer, error) {
	dialerStats.ConnReuse.Add(1)
	conn.lastAccessTime = time.Now()
	conn.refs++
	dialer.evictSorted = false
	return dialer.connWithCloser(conn)
}

// connWithCloser returns the three-tuple expected by the main dial func, where
// the closer handles the correct state management for updating the conns place
// in the eviction queue.
func (dialer *cachedConnDialer) connWithCloser(conn *cachedConn) (tabletmanagerservicepb.TabletManagerClient, io.Closer, error) {
	return conn, closeFunc(func() error {
		dialer.m.Lock()
		defer dialer.m.Unlock()
		conn.refs--
		dialer.evictSorted = false
		return nil
	}), nil
}

// Close closes all currently cached connections, ***regardless of whether
// those connections are in use***. Calling Close therefore will fail any RPCs
// using currently lent-out connections, and, furthermore, will invalidate the
// io.Closer that was returned for that connection from dialer.dial().
//
// As a result, it is not safe to reuse a cachedConnDialer after calling Close,
// and you should instead obtain a new one by calling either
// tmclient.TabletManagerClient() with
// TabletManagerProtocol set to "grpc-cached", or by calling
// grpctmclient.NewCachedConnClient directly.
func (dialer *cachedConnDialer) Close() {
	dialer.m.Lock()
	defer dialer.m.Unlock()

	for _, conn := range dialer.evict {
		conn.cc.Close()
		delete(dialer.conns, conn.addr)
		dialer.connWaitSema.Release()
	}
	dialer.evict = nil
}

func getTabletAddr(tablet *topodatapb.Tablet) string {
	return netutil.JoinHostPort(tablet.Hostname, int32(tablet.PortMap["grpc"]))
}
<commit_msg>Add missing connwaitsema release on a not-actually-new newdial<commit_after>/*
Copyright 2021 The Vitess 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 grpctmclient

import (
	"context"
	"flag"
	"io"
	"sort"
	"sync"
	"time"

	"google.golang.org/grpc"

	"vitess.io/vitess/go/netutil"
	"vitess.io/vitess/go/stats"
	"vitess.io/vitess/go/sync2"
	"vitess.io/vitess/go/vt/grpcclient"
	"vitess.io/vitess/go/vt/vttablet/tmclient"

	tabletmanagerservicepb "vitess.io/vitess/go/vt/proto/tabletmanagerservice"
	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
)

var (
	defaultPoolCapacity = flag.Int("tablet_manager_grpc_connpool_size", 100, "number of tablets to keep tmclient connections open to")
)

func init() {
	tmclient.RegisterTabletManagerClientFactory("grpc-cached", func() tmclient.TabletManagerClient {
		return NewCachedConnClient(*defaultPoolCapacity)
	})
}

// closeFunc allows a standalone function to implement io.Closer, similar to
// how http.HandlerFunc allows standalone functions to implement http.Handler.
type closeFunc func() error

func (fn closeFunc) Close() error {
	return fn()
}

var _ io.Closer = (*closeFunc)(nil)

type cachedConn struct {
	tabletmanagerservicepb.TabletManagerClient
	cc *grpc.ClientConn

	addr           string
	lastAccessTime time.Time
	refs           int
}

type cachedConnDialer struct {
	m            sync.Mutex
	conns        map[string]*cachedConn
	evict        []*cachedConn
	evictSorted  bool
	connWaitSema *sync2.Semaphore
}

var dialerStats = struct {
	ConnReuse    *stats.Gauge
	ConnNew      *stats.Gauge
	DialTimeouts *stats.Gauge
	DialTimings  *stats.Timings
}{
	ConnReuse:    stats.NewGauge("tabletmanagerclient_cachedconn_reuse", "number of times a call to dial() was able to reuse an existing connection"),
	ConnNew:      stats.NewGauge("tabletmanagerclient_cachedconn_new", "number of times a call to dial() resulted in a dialing a new grpc clientconn"),
	DialTimeouts: stats.NewGauge("tabletmanagerclient_cachedconn_dial_timeouts", "number of context timeouts during dial()"),
	DialTimings:  stats.NewTimings("tabletmanagerclient_cachedconn_dialtimings", "timings for various dial paths", "path", "rlock_fast", "sema_fast", "sema_poll"),
}

// NewCachedConnClient returns a grpc Client that caches connections to the different tablets
func NewCachedConnClient(capacity int) *Client {
	dialer := &cachedConnDialer{
		conns:        make(map[string]*cachedConn, capacity),
		evict:        make([]*cachedConn, 0, capacity),
		connWaitSema: sync2.NewSemaphore(capacity, 0),
	}
	return &Client{dialer}
}

var _ dialer = (*cachedConnDialer)(nil)

func (dialer *cachedConnDialer) sortEvictionsLocked() {
	if !dialer.evictSorted {
		sort.Slice(dialer.evict, func(i, j int) bool {
			left, right := dialer.evict[i], dialer.evict[j]
			if left.refs == right.refs {
				return left.lastAccessTime.After(right.lastAccessTime)
			}
			return left.refs > right.refs
		})
		dialer.evictSorted = true
	}
}

func (dialer *cachedConnDialer) dial(ctx context.Context, tablet *topodatapb.Tablet) (tabletmanagerservicepb.TabletManagerClient, io.Closer, error) {
	start := time.Now()
	addr := getTabletAddr(tablet)

	if client, closer, found, err := dialer.tryFromCache(addr, &dialer.m); found {
		dialerStats.DialTimings.Add("rlock_fast", time.Since(start))
		return client, closer, err
	}

	if dialer.connWaitSema.TryAcquire() {
		defer func() {
			dialerStats.DialTimings.Add("sema_fast", time.Since(start))
		}()

		// Check if another goroutine managed to dial a conn for the same addr
		// while we were waiting for the write lock. This is identical to the
		// read-lock section above.
		if client, closer, found, err := dialer.tryFromCache(addr, &dialer.m); found {
			return client, closer, err
		}
		return dialer.newdial(addr)
	}

	defer func() {
		dialerStats.DialTimings.Add("sema_poll", time.Since(start))
	}()

	for {
		select {
		case <-ctx.Done():
			dialerStats.DialTimeouts.Add(1)
			return nil, nil, ctx.Err()
		default:
			if client, closer, found, err := dialer.pollOnce(addr); found {
				return client, closer, err
			}
		}
	}
}

// tryFromCache tries to get a connection from the cache, performing a redial
// on that connection if it exists. It returns a TabletManagerClient impl, an
// io.Closer, a flag to indicate whether a connection was found in the cache,
// and an error, which is always nil.
//
// In addition to the addr being dialed, tryFromCache takes a sync.Locker which,
// if not nil, will be used to wrap the lookup and redial in that lock. This
// function can be called in situations where the conns map is locked
// externally (like in pollOnce), so we do not want to manage the locks here. In
// other cases (like in the rlock_fast path of dial()), we pass in the RLocker
// to ensure we have a read lock on the cache for the duration of the call.
func (dialer *cachedConnDialer) tryFromCache(addr string, locker sync.Locker) (client tabletmanagerservicepb.TabletManagerClient, closer io.Closer, found bool, err error) {
	if locker != nil {
		locker.Lock()
		defer locker.Unlock()
	}

	if conn, ok := dialer.conns[addr]; ok {
		client, closer, err := dialer.redialLocked(conn)
		return client, closer, ok, err
	}

	return nil, nil, false, nil
}

// pollOnce is called on each iteration of the polling loop in dial(). It:
// - locks the conns cache for writes
// - attempts to get a connection from the cache. If found, redial() it and exit.
// - locks the queue
// - peeks at the head of the eviction queue. if the peeked conn has no refs, it
//   is unused, and can be evicted to make room for the new connection to addr.
//   If the peeked conn has refs, exit.
// - pops the conn we just peeked from the queue, delete it from the cache, and
//   close the underlying ClientConn for that conn.
// - attempt a newdial. if the newdial fails, it will release a slot on the
//   connWaitSema, so another dial() call can successfully acquire it to dial
//   a new conn. if the newdial succeeds, we will have evicted one conn, but
//   added another, so the net change is 0, and no changes to the connWaitSema
//   are made.
//
// It returns a TabletManagerClient impl, an io.Closer, a flag to indicate
// whether the dial() poll loop should exit, and an error.
func (dialer *cachedConnDialer) pollOnce(addr string) (client tabletmanagerservicepb.TabletManagerClient, closer io.Closer, found bool, err error) {
	dialer.m.Lock()

	if client, closer, found, err := dialer.tryFromCache(addr, nil); found {
		dialer.m.Unlock()
		return client, closer, found, err
	}

	dialer.sortEvictionsLocked()

	conn := dialer.evict[len(dialer.evict)-1]
	if conn.refs != 0 {
		dialer.m.Unlock()
		return nil, nil, false, nil
	}

	dialer.evict = dialer.evict[:len(dialer.evict)-1]
	delete(dialer.conns, conn.addr)
	conn.cc.Close()
	dialer.m.Unlock()

	client, closer, err = dialer.newdial(addr)
	return client, closer, true, err
}

// newdial creates a new cached connection, and updates the cache and eviction
// queue accordingly. If newdial fails to create the underlying
// gRPC connection, it will make a call to Release the connWaitSema for other
// newdial calls.
//
// It returns the three-tuple of client-interface, closer, and error that the
// main dial func returns.
func (dialer *cachedConnDialer) newdial(addr string) (tabletmanagerservicepb.TabletManagerClient, io.Closer, error) {
	dialerStats.ConnNew.Add(1)

	opt, err := grpcclient.SecureDialOption(*cert, *key, *ca, *name)
	if err != nil {
		dialer.connWaitSema.Release()
		return nil, nil, err
	}

	cc, err := grpcclient.Dial(addr, grpcclient.FailFast(false), opt)
	if err != nil {
		dialer.connWaitSema.Release()
		return nil, nil, err
	}

	dialer.m.Lock()
	defer dialer.m.Unlock()

	if conn, existing := dialer.conns[addr]; existing {
		// race condition: some other goroutine has dialed our tablet before we have;
		// this is not great, but shouldn't happen often (if at all), so we're going to
		// close this connection and reuse the existing one. by doing this, we can keep
		// the actual Dial out of the global lock and significantly increase throughput
		cc.Close()
		dialer.connWaitSema.Release()
		return dialer.redialLocked(conn)
	}

	conn := &cachedConn{
		TabletManagerClient: tabletmanagerservicepb.NewTabletManagerClient(cc),
		cc:                  cc,
		lastAccessTime:      time.Now(),
		refs:                1,
		addr:                addr,
	}
	dialer.evict = append(dialer.evict, conn)
	dialer.evictSorted = false
	dialer.conns[addr] = conn

	return dialer.connWithCloser(conn)
}

// redialLocked takes an already-dialed connection in the cache does all the work of
// lending that connection out to one more caller. this should only ever be
// called while holding at least the RLock on dialer.m (but the write lock is
// fine too), to prevent the connection from getting evicted out from under us.
//
// It returns the three-tuple of client-interface, closer, and error that the
// main dial func returns.
func (dialer *cachedConnDialer) redialLocked(conn *cachedConn) (tabletmanagerservicepb.TabletManagerClient, io.Closer, error) {
	dialerStats.ConnReuse.Add(1)
	conn.lastAccessTime = time.Now()
	conn.refs++
	dialer.evictSorted = false
	return dialer.connWithCloser(conn)
}

// connWithCloser returns the three-tuple expected by the main dial func, where
// the closer handles the correct state management for updating the conns place
// in the eviction queue.
func (dialer *cachedConnDialer) connWithCloser(conn *cachedConn) (tabletmanagerservicepb.TabletManagerClient, io.Closer, error) {
	return conn, closeFunc(func() error {
		dialer.m.Lock()
		defer dialer.m.Unlock()
		conn.refs--
		dialer.evictSorted = false
		return nil
	}), nil
}

// Close closes all currently cached connections, ***regardless of whether
// those connections are in use***. Calling Close therefore will fail any RPCs
// using currently lent-out connections, and, furthermore, will invalidate the
// io.Closer that was returned for that connection from dialer.dial().
//
// As a result, it is not safe to reuse a cachedConnDialer after calling Close,
// and you should instead obtain a new one by calling either
// tmclient.TabletManagerClient() with
// TabletManagerProtocol set to "grpc-cached", or by calling
// grpctmclient.NewCachedConnClient directly.
func (dialer *cachedConnDialer) Close() {
	dialer.m.Lock()
	defer dialer.m.Unlock()

	for _, conn := range dialer.evict {
		conn.cc.Close()
		delete(dialer.conns, conn.addr)
		dialer.connWaitSema.Release()
	}
	dialer.evict = nil
}

func getTabletAddr(tablet *topodatapb.Tablet) string {
	return netutil.JoinHostPort(tablet.Hostname, int32(tablet.PortMap["grpc"]))
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess 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 connpool

import (
	"net"
	"sync"
	"time"

	"vitess.io/vitess/go/netutil"

	"context"

	"vitess.io/vitess/go/pools"
	"vitess.io/vitess/go/sync2"
	"vitess.io/vitess/go/trace"
	"vitess.io/vitess/go/vt/callerid"
	"vitess.io/vitess/go/vt/dbconfigs"
	"vitess.io/vitess/go/vt/dbconnpool"
	"vitess.io/vitess/go/vt/log"
	"vitess.io/vitess/go/vt/mysqlctl"
	"vitess.io/vitess/go/vt/vterrors"
	"vitess.io/vitess/go/vt/vttablet/tabletserver/tabletenv"

	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
)

// ErrConnPoolClosed is returned when the connection pool is closed.
var ErrConnPoolClosed = vterrors.New(vtrpcpb.Code_INTERNAL, "internal error: unexpected: conn pool is closed")

// Pool implements a custom connection pool for tabletserver.
// It's similar to dbconnpool.ConnPool, but the connections it creates
// come with built-in ability to kill in-flight queries. These connections
// also trigger a CheckMySQL call if we fail to connect to MySQL.
// Other than the connection type, ConnPool maintains an additional
// pool of dba connections that are used to kill connections.
type Pool struct {
	env                tabletenv.Env
	name               string
	mu                 sync.Mutex
	connections        *pools.ResourcePool
	capacity           int
	prefillParallelism int
	timeout            time.Duration
	idleTimeout        time.Duration
	waiterCap          int64
	waiterCount        sync2.AtomicInt64
	dbaPool            *dbconnpool.ConnectionPool
	appDebugParams     dbconfigs.Connector
}

// NewPool creates a new Pool. The name is used
// to publish stats only.
func NewPool(env tabletenv.Env, name string, cfg tabletenv.ConnPoolConfig) *Pool {
	idleTimeout := cfg.IdleTimeoutSeconds.Get()
	cp := &Pool{
		env:                env,
		name:               name,
		capacity:           cfg.Size,
		prefillParallelism: cfg.PrefillParallelism,
		timeout:            cfg.TimeoutSeconds.Get(),
		idleTimeout:        idleTimeout,
		waiterCap:          int64(cfg.MaxWaiters),
		dbaPool:            dbconnpool.NewConnectionPool("", 1, idleTimeout, 0),
	}
	if name == "" {
		return cp
	}
	env.Exporter().NewGaugeFunc(name+"Capacity", "Tablet server conn pool capacity", cp.Capacity)
	env.Exporter().NewGaugeFunc(name+"Available", "Tablet server conn pool available", cp.Available)
	env.Exporter().NewGaugeFunc(name+"Active", "Tablet server conn pool active", cp.Active)
	env.Exporter().NewGaugeFunc(name+"InUse", "Tablet server conn pool in use", cp.InUse)
	env.Exporter().NewGaugeFunc(name+"MaxCap", "Tablet server conn pool max cap", cp.MaxCap)
	env.Exporter().NewCounterFunc(name+"WaitCount", "Tablet server conn pool wait count", cp.WaitCount)
	env.Exporter().NewCounterDurationFunc(name+"WaitTime", "Tablet server wait time", cp.WaitTime)
	env.Exporter().NewGaugeDurationFunc(name+"IdleTimeout", "Tablet server idle timeout", cp.IdleTimeout)
	env.Exporter().NewCounterFunc(name+"IdleClosed", "Tablet server conn pool idle closed", cp.IdleClosed)
	env.Exporter().NewCounterFunc(name+"Exhausted", "Number of times pool had zero available slots", cp.Exhausted)
	return cp
}

func (cp *Pool) pool() (p *pools.ResourcePool) {
	cp.mu.Lock()
	p = cp.connections
	cp.mu.Unlock()
	return p
}

// Open must be called before starting to use the pool.
func (cp *Pool) Open(appParams, dbaParams, appDebugParams dbconfigs.Connector) {
	cp.mu.Lock()
	defer cp.mu.Unlock()

	if cp.prefillParallelism != 0 {
		log.Infof("Opening pool: '%s'", cp.name)
		defer log.Infof("Done opening pool: '%s'", cp.name)
	}

	f := func(ctx context.Context) (pools.Resource, error) {
		return NewDBConn(ctx, cp, appParams)
	}

	var refreshCheck pools.RefreshCheck
	if net.ParseIP(appParams.Host()) == nil {
		refreshCheck = netutil.DNSTracker(appParams.Host())
	}

	cp.connections = pools.NewResourcePool(f, cp.capacity, cp.capacity, cp.idleTimeout, cp.prefillParallelism, cp.getLogWaitCallback(), refreshCheck, *mysqlctl.PoolDynamicHostnameResolution)
	cp.appDebugParams = appDebugParams

	cp.dbaPool.Open(dbaParams)
}

func (cp *Pool) getLogWaitCallback() func(time.Time) {
	if cp.name == "" {
		return func(start time.Time) {} // no op
	}
	return func(start time.Time) {
		cp.env.Stats().WaitTimings.Record(cp.name+"ResourceWaitTime", start)
	}
}

// Close will close the pool and wait for connections to be returned before
// exiting.
func (cp *Pool) Close() {
	p := cp.pool()
	if p == nil {
		return
	}
	// We should not hold the lock while calling Close
	// because it waits for connections to be returned.
	p.Close()
	cp.mu.Lock()
	cp.connections = nil
	cp.mu.Unlock()
	cp.dbaPool.Close()
}

// Get returns a connection.
// You must call Recycle on DBConn once done.
func (cp *Pool) Get(ctx context.Context) (*DBConn, error) {
	span, ctx := trace.NewSpan(ctx, "Pool.Get")
	defer span.Finish()

	if cp.waiterCap > 0 {
		waiterCount := cp.waiterCount.Add(1)
		defer cp.waiterCount.Add(-1)
		if waiterCount > cp.waiterCap {
			return nil, vterrors.Errorf(vtrpcpb.Code_RESOURCE_EXHAUSTED, "pool %s waiter count exceeded", cp.name)
		}
	}

	if cp.isCallerIDAppDebug(ctx) {
		return NewDBConnNoPool(ctx, cp.appDebugParams, cp.dbaPool)
	}
	p := cp.pool()
	if p == nil {
		return nil, ErrConnPoolClosed
	}
	span.Annotate("capacity", p.Capacity())
	span.Annotate("in_use", p.InUse())
	span.Annotate("available", p.Available())
	span.Annotate("active", p.Active())

	if cp.timeout != 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, cp.timeout)
		defer cancel()
	}
	r, err := p.Get(ctx)
	if err != nil {
		return nil, err
	}
	return r.(*DBConn), nil
}

// Put puts a connection into the pool.
func (cp *Pool) Put(conn *DBConn) {
	p := cp.pool()
	if p == nil {
		panic(ErrConnPoolClosed)
	}
	if conn == nil {
		p.Put(nil)
	} else {
		p.Put(conn)
	}
}

// SetCapacity alters the size of the pool at runtime.
func (cp *Pool) SetCapacity(capacity int) (err error) {
	cp.mu.Lock()
	defer cp.mu.Unlock()
	if cp.connections != nil {
		err = cp.connections.SetCapacity(capacity)
		if err != nil {
			return err
		}
	}
	cp.capacity = capacity
	return nil
}

// SetIdleTimeout sets the idleTimeout on the pool.
func (cp *Pool) SetIdleTimeout(idleTimeout time.Duration) {
	cp.mu.Lock()
	defer cp.mu.Unlock()
	if cp.connections != nil {
		cp.connections.SetIdleTimeout(idleTimeout)
	}
	cp.dbaPool.SetIdleTimeout(idleTimeout)
	cp.idleTimeout = idleTimeout
}

// StatsJSON returns the pool stats as a JSON object.
func (cp *Pool) StatsJSON() string {
	p := cp.pool()
	if p == nil {
		return "{}"
	}
	return p.StatsJSON()
}

// Capacity returns the pool capacity.
func (cp *Pool) Capacity() int64 {
	p := cp.pool()
	if p == nil {
		return 0
	}
	return p.Capacity()
}

// Available returns the number of available connections in the pool
func (cp *Pool) Available() int64 {
	p := cp.pool()
	if p == nil {
		return 0
	}
	return p.Available()
}

// Active returns the number of active connections in the pool
func (cp *Pool) Active() int64 {
	p := cp.pool()
	if p == nil {
		return 0
	}
	return p.Active()
}

// InUse returns the number of in-use connections in the pool
func (cp *Pool) InUse() int64 {
	p := cp.pool()
	if p == nil {
		return 0
	}
	return p.InUse()
}

// MaxCap returns the maximum size of the pool
func (cp *Pool) MaxCap() int64 {
	p := cp.pool()
	if p == nil {
		return 0
	}
	return p.MaxCap()
}

// WaitCount returns how many clients are waiting for a connection
func (cp *Pool) WaitCount() int64 {
	p := cp.pool()
	if p == nil {
		return 0
	}
	return p.WaitCount()
}

// WaitTime return the pool WaitTime.
func (cp *Pool) WaitTime() time.Duration {
	p := cp.pool()
	if p == nil {
		return 0
	}
	return p.WaitTime()
}

// IdleTimeout returns the idle timeout for the pool.
func (cp *Pool) IdleTimeout() time.Duration {
	p := cp.pool()
	if p == nil {
		return 0
	}
	return p.IdleTimeout()
}

// IdleClosed returns the number of closed connections for the pool.
func (cp *Pool) IdleClosed() int64 {
	p := cp.pool()
	if p == nil {
		return 0
	}
	return p.IdleClosed()
}

// Exhausted returns the number of times available went to zero for the pool.
func (cp *Pool) Exhausted() int64 {
	p := cp.pool()
	if p == nil {
		return 0
	}
	return p.Exhausted()
}

func (cp *Pool) isCallerIDAppDebug(ctx context.Context) bool {
	params, err := cp.appDebugParams.MysqlParams()
	if err != nil {
		return false
	}
	if params == nil || params.Uname == "" {
		return false
	}
	callerID := callerid.ImmediateCallerIDFromContext(ctx)
	return callerID != nil && callerID.Username == params.Uname
}
<commit_msg>Add stats to track when a pool's waiter queue is full<commit_after>/*
Copyright 2019 The Vitess 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 connpool

import (
	"net"
	"sync"
	"time"

	"vitess.io/vitess/go/netutil"

	"context"

	"vitess.io/vitess/go/pools"
	"vitess.io/vitess/go/sync2"
	"vitess.io/vitess/go/trace"
	"vitess.io/vitess/go/vt/callerid"
	"vitess.io/vitess/go/vt/dbconfigs"
	"vitess.io/vitess/go/vt/dbconnpool"
	"vitess.io/vitess/go/vt/log"
	"vitess.io/vitess/go/vt/mysqlctl"
	"vitess.io/vitess/go/vt/vterrors"
	"vitess.io/vitess/go/vt/vttablet/tabletserver/tabletenv"

	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
)

// ErrConnPoolClosed is returned when the connection pool is closed.
var ErrConnPoolClosed = vterrors.New(vtrpcpb.Code_INTERNAL, "internal error: unexpected: conn pool is closed")

// Pool implements a custom connection pool for tabletserver.
// It's similar to dbconnpool.ConnPool, but the connections it creates
// come with built-in ability to kill in-flight queries. These connections
// also trigger a CheckMySQL call if we fail to connect to MySQL.
// Other than the connection type, ConnPool maintains an additional
// pool of dba connections that are used to kill connections.
type Pool struct {
	env                tabletenv.Env
	name               string
	mu                 sync.Mutex
	connections        *pools.ResourcePool
	capacity           int
	prefillParallelism int
	timeout            time.Duration
	idleTimeout        time.Duration
	waiterCap          int64
	waiterCount        sync2.AtomicInt64
	waiterQueueFull    sync2.AtomicInt64
	dbaPool            *dbconnpool.ConnectionPool
	appDebugParams     dbconfigs.Connector
}

// NewPool creates a new Pool. The name is used
// to publish stats only.
func NewPool(env tabletenv.Env, name string, cfg tabletenv.ConnPoolConfig) *Pool {
	idleTimeout := cfg.IdleTimeoutSeconds.Get()
	cp := &Pool{
		env:                env,
		name:               name,
		capacity:           cfg.Size,
		prefillParallelism: cfg.PrefillParallelism,
		timeout:            cfg.TimeoutSeconds.Get(),
		idleTimeout:        idleTimeout,
		waiterCap:          int64(cfg.MaxWaiters),
		dbaPool:            dbconnpool.NewConnectionPool("", 1, idleTimeout, 0),
	}
	if name == "" {
		return cp
	}
	env.Exporter().NewGaugeFunc(name+"Capacity", "Tablet server conn pool capacity", cp.Capacity)
	env.Exporter().NewGaugeFunc(name+"Available", "Tablet server conn pool available", cp.Available)
	env.Exporter().NewGaugeFunc(name+"Active", "Tablet server conn pool active", cp.Active)
	env.Exporter().NewGaugeFunc(name+"InUse", "Tablet server conn pool in use", cp.InUse)
	env.Exporter().NewGaugeFunc(name+"MaxCap", "Tablet server conn pool max cap", cp.MaxCap)
	env.Exporter().NewCounterFunc(name+"WaitCount", "Tablet server conn pool wait count", cp.WaitCount)
	env.Exporter().NewCounterDurationFunc(name+"WaitTime", "Tablet server wait time", cp.WaitTime)
	env.Exporter().NewGaugeDurationFunc(name+"IdleTimeout", "Tablet server idle timeout", cp.IdleTimeout)
	env.Exporter().NewCounterFunc(name+"IdleClosed", "Tablet server conn pool idle closed", cp.IdleClosed)
	env.Exporter().NewCounterFunc(name+"Exhausted", "Number of times pool had zero available slots", cp.Exhausted)
	env.Exporter().NewCounterFunc(name+"WaiterQueueFull", "Number of times the waiter queue was full", cp.waiterQueueFull.Get)
	return cp
}

func (cp *Pool) pool() (p *pools.ResourcePool) {
	cp.mu.Lock()
	p = cp.connections
	cp.mu.Unlock()
	return p
}

// Open must be called before starting to use the pool.
func (cp *Pool) Open(appParams, dbaParams, appDebugParams dbconfigs.Connector) {
	cp.mu.Lock()
	defer cp.mu.Unlock()

	if cp.prefillParallelism != 0 {
		log.Infof("Opening pool: '%s'", cp.name)
		defer log.Infof("Done opening pool: '%s'", cp.name)
	}

	f := func(ctx context.Context) (pools.Resource, error) {
		return NewDBConn(ctx, cp, appParams)
	}

	var refreshCheck pools.RefreshCheck
	if net.ParseIP(appParams.Host()) == nil {
		refreshCheck = netutil.DNSTracker(appParams.Host())
	}

	cp.connections = pools.NewResourcePool(f, cp.capacity, cp.capacity, cp.idleTimeout, cp.prefillParallelism, cp.getLogWaitCallback(), refreshCheck, *mysqlctl.PoolDynamicHostnameResolution)
	cp.appDebugParams = appDebugParams

	cp.dbaPool.Open(dbaParams)
}

func (cp *Pool) getLogWaitCallback() func(time.Time) {
	if cp.name == "" {
		return func(start time.Time) {} // no op
	}
	return func(start time.Time) {
		cp.env.Stats().WaitTimings.Record(cp.name+"ResourceWaitTime", start)
	}
}

// Close will close the pool and wait for connections to be returned before
// exiting.
func (cp *Pool) Close() {
	p := cp.pool()
	if p == nil {
		return
	}
	// We should not hold the lock while calling Close
	// because it waits for connections to be returned.
	p.Close()
	cp.mu.Lock()
	cp.connections = nil
	cp.mu.Unlock()
	cp.dbaPool.Close()
}

// Get returns a connection.
// You must call Recycle on DBConn once done.
func (cp *Pool) Get(ctx context.Context) (*DBConn, error) {
	span, ctx := trace.NewSpan(ctx, "Pool.Get")
	defer span.Finish()

	if cp.waiterCap > 0 {
		waiterCount := cp.waiterCount.Add(1)
		defer cp.waiterCount.Add(-1)
		if waiterCount > cp.waiterCap {
			cp.waiterQueueFull.Add(1)
			return nil, vterrors.Errorf(vtrpcpb.Code_RESOURCE_EXHAUSTED, "pool %s waiter count exceeded", cp.name)
		}
	}

	if cp.isCallerIDAppDebug(ctx) {
		return NewDBConnNoPool(ctx, cp.appDebugParams, cp.dbaPool)
	}
	p := cp.pool()
	if p == nil {
		return nil, ErrConnPoolClosed
	}
	span.Annotate("capacity", p.Capacity())
	span.Annotate("in_use", p.InUse())
	span.Annotate("available", p.Available())
	span.Annotate("active", p.Active())

	if cp.timeout != 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, cp.timeout)
		defer cancel()
	}
	r, err := p.Get(ctx)
	if err != nil {
		return nil, err
	}
	return r.(*DBConn), nil
}

// Put puts a connection into the pool.
func (cp *Pool) Put(conn *DBConn) {
	p := cp.pool()
	if p == nil {
		panic(ErrConnPoolClosed)
	}
	if conn == nil {
		p.Put(nil)
	} else {
		p.Put(conn)
	}
}

// SetCapacity alters the size of the pool at runtime.
func (cp *Pool) SetCapacity(capacity int) (err error) {
	cp.mu.Lock()
	defer cp.mu.Unlock()
	if cp.connections != nil {
		err = cp.connections.SetCapacity(capacity)
		if err != nil {
			return err
		}
	}
	cp.capacity = capacity
	return nil
}

// SetIdleTimeout sets the idleTimeout on the pool.
func (cp *Pool) SetIdleTimeout(idleTimeout time.Duration) {
	cp.mu.Lock()
	defer cp.mu.Unlock()
	if cp.connections != nil {
		cp.connections.SetIdleTimeout(idleTimeout)
	}
	cp.dbaPool.SetIdleTimeout(idleTimeout)
	cp.idleTimeout = idleTimeout
}

// StatsJSON returns the pool stats as a JSON object.
func (cp *Pool) StatsJSON() string {
	p := cp.pool()
	if p == nil {
		return "{}"
	}
	return p.StatsJSON()
}

// Capacity returns the pool capacity.
func (cp *Pool) Capacity() int64 {
	p := cp.pool()
	if p == nil {
		return 0
	}
	return p.Capacity()
}

// Available returns the number of available connections in the pool
func (cp *Pool) Available() int64 {
	p := cp.pool()
	if p == nil {
		return 0
	}
	return p.Available()
}

// Active returns the number of active connections in the pool
func (cp *Pool) Active() int64 {
	p := cp.pool()
	if p == nil {
		return 0
	}
	return p.Active()
}

// InUse returns the number of in-use connections in the pool
func (cp *Pool) InUse() int64 {
	p := cp.pool()
	if p == nil {
		return 0
	}
	return p.InUse()
}

// MaxCap returns the maximum size of the pool
func (cp *Pool) MaxCap() int64 {
	p := cp.pool()
	if p == nil {
		return 0
	}
	return p.MaxCap()
}

// WaitCount returns how many clients are waiting for a connection
func (cp *Pool) WaitCount() int64 {
	p := cp.pool()
	if p == nil {
		return 0
	}
	return p.WaitCount()
}

// WaitTime return the pool WaitTime.
func (cp *Pool) WaitTime() time.Duration {
	p := cp.pool()
	if p == nil {
		return 0
	}
	return p.WaitTime()
}

// IdleTimeout returns the idle timeout for the pool.
func (cp *Pool) IdleTimeout() time.Duration {
	p := cp.pool()
	if p == nil {
		return 0
	}
	return p.IdleTimeout()
}

// IdleClosed returns the number of closed connections for the pool.
func (cp *Pool) IdleClosed() int64 {
	p := cp.pool()
	if p == nil {
		return 0
	}
	return p.IdleClosed()
}

// Exhausted returns the number of times available went to zero for the pool.
func (cp *Pool) Exhausted() int64 {
	p := cp.pool()
	if p == nil {
		return 0
	}
	return p.Exhausted()
}

func (cp *Pool) isCallerIDAppDebug(ctx context.Context) bool {
	params, err := cp.appDebugParams.MysqlParams()
	if err != nil {
		return false
	}
	if params == nil || params.Uname == "" {
		return false
	}
	callerID := callerid.ImmediateCallerIDFromContext(ctx)
	return callerID != nil && callerID.Username == params.Uname
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
// Author: jacobsa@google.com (Aaron Jacobs)

package gcs

import (
	"errors"
	"io"
	"net/http"
	"unicode/utf8"

	"golang.org/x/net/context"
	"google.golang.org/cloud"
	"google.golang.org/cloud/storage"
)

// A request to create an object, accepted by Bucket.CreateObject.
type CreateObjectRequest struct {
	// Attributes with which the object should be created. The Name field must be
	// set; other zero-valued fields are ignored.
	//
	// Object names must:
	//
	// *  be non-empty.
	// *  be no longer than 1024 bytes.
	// *  be valid UTF-8.
	// *  not contain the code point U+000A (line feed).
	// *  not contain the code point U+000D (carriage return).
	//
	// See here for authoritative documentation:
	//     https://cloud.google.com/storage/docs/bucket-naming#objectnames
	Attrs storage.ObjectAttrs

	// A reader from which to obtain the contents of the object. Must be non-nil.
	// Will be closed when CreateObject returns, regardless of whether it succeeds.
	Contents io.ReadCloser
}

// Bucket represents a GCS bucket, pre-bound with a bucket name and necessary
// authorization information.
//
// Each method that may block accepts a context object that is used for
// deadlines and cancellation. Users need not package authorization information
// into the context object (using cloud.WithContext or similar).
type Bucket interface {
	Name() string

	// List the objects in the bucket that meet the criteria defined by the
	// query, returning a result object that contains the results and potentially
	// a cursor for retrieving the next portion of the larger set of results.
	ListObjects(ctx context.Context, query *storage.Query) (*storage.Objects, error)

	// Create a reader for the contents of the object with the given name. The
	// caller must arrange for the reader to be closed when it is no longer
	// needed.
	NewReader(ctx context.Context, objectName string) (io.ReadCloser, error)

	// Create or overwrite an object according to the supplied request. The new
	// object is guaranteed to exist immediately for the purposes of reading (and
	// eventually for listing) after this method returns a nil error. It is
	// guaranteed not to exist before req.Contents returns io.EOF.
	CreateObject(ctx context.Context, req *CreateObjectRequest) (*storage.Object, error)

	// Delete the object with the given name.
	DeleteObject(ctx context.Context, name string) error
}

type bucket struct {
	projID string
	client *http.Client
	name   string
}

func (b *bucket) Name() string {
	return b.name
}

func (b *bucket) ListObjects(ctx context.Context, query *storage.Query) (*storage.Objects, error) {
	authContext := cloud.WithContext(ctx, b.projID, b.client)
	return storage.ListObjects(authContext, b.name, query)
}

func (b *bucket) NewReader(ctx context.Context, objectName string) (io.ReadCloser, error) {
	authContext := cloud.WithContext(ctx, b.projID, b.client)
	return storage.NewReader(authContext, b.name, objectName)
}

func (b *bucket) NewWriter(
	ctx context.Context,
	attrs *storage.ObjectAttrs) (ObjectWriter, error) {
	authContext := cloud.WithContext(ctx, b.projID, b.client)

	// As of 2015-02, the wrapped storage package doesn't check this for us,
	// causing silently transformed names:
	//     https://github.com/GoogleCloudPlatform/gcloud-golang/issues/111
	if !utf8.ValidString(attrs.Name) {
		return nil, errors.New("Invalid object name: not valid UTF-8")
	}

	// Create and initialize the wrapped writer.
	wrapped := storage.NewWriter(authContext, b.name, attrs.Name)
	wrapped.ObjectAttrs = *attrs

	w := &objectWriter{
		wrapped: wrapped,
	}

	return w, nil
}

func (b *bucket) DeleteObject(ctx context.Context, name string) error {
	authContext := cloud.WithContext(ctx, b.projID, b.client)
	return storage.DeleteObject(authContext, b.name, name)
}
<commit_msg>Contents doesn't need to be a ReadCloser, I think.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
// Author: jacobsa@google.com (Aaron Jacobs)

package gcs

import (
	"errors"
	"io"
	"net/http"
	"unicode/utf8"

	"golang.org/x/net/context"
	"google.golang.org/cloud"
	"google.golang.org/cloud/storage"
)

// A request to create an object, accepted by Bucket.CreateObject.
type CreateObjectRequest struct {
	// Attributes with which the object should be created. The Name field must be
	// set; other zero-valued fields are ignored.
	//
	// Object names must:
	//
	// *  be non-empty.
	// *  be no longer than 1024 bytes.
	// *  be valid UTF-8.
	// *  not contain the code point U+000A (line feed).
	// *  not contain the code point U+000D (carriage return).
	//
	// See here for authoritative documentation:
	//     https://cloud.google.com/storage/docs/bucket-naming#objectnames
	Attrs storage.ObjectAttrs

	// A reader from which to obtain the contents of the object. Must be non-nil.
	Contents io.Reader
}

// Bucket represents a GCS bucket, pre-bound with a bucket name and necessary
// authorization information.
//
// Each method that may block accepts a context object that is used for
// deadlines and cancellation. Users need not package authorization information
// into the context object (using cloud.WithContext or similar).
type Bucket interface {
	Name() string

	// List the objects in the bucket that meet the criteria defined by the
	// query, returning a result object that contains the results and potentially
	// a cursor for retrieving the next portion of the larger set of results.
	ListObjects(ctx context.Context, query *storage.Query) (*storage.Objects, error)

	// Create a reader for the contents of the object with the given name. The
	// caller must arrange for the reader to be closed when it is no longer
	// needed.
	NewReader(ctx context.Context, objectName string) (io.ReadCloser, error)

	// Create or overwrite an object according to the supplied request. The new
	// object is guaranteed to exist immediately for the purposes of reading (and
	// eventually for listing) after this method returns a nil error. It is
	// guaranteed not to exist before req.Contents returns io.EOF.
	CreateObject(ctx context.Context, req *CreateObjectRequest) (*storage.Object, error)

	// Delete the object with the given name.
	DeleteObject(ctx context.Context, name string) error
}

type bucket struct {
	projID string
	client *http.Client
	name   string
}

func (b *bucket) Name() string {
	return b.name
}

func (b *bucket) ListObjects(ctx context.Context, query *storage.Query) (*storage.Objects, error) {
	authContext := cloud.WithContext(ctx, b.projID, b.client)
	return storage.ListObjects(authContext, b.name, query)
}

func (b *bucket) NewReader(ctx context.Context, objectName string) (io.ReadCloser, error) {
	authContext := cloud.WithContext(ctx, b.projID, b.client)
	return storage.NewReader(authContext, b.name, objectName)
}

func (b *bucket) NewWriter(
	ctx context.Context,
	attrs *storage.ObjectAttrs) (ObjectWriter, error) {
	authContext := cloud.WithContext(ctx, b.projID, b.client)

	// As of 2015-02, the wrapped storage package doesn't check this for us,
	// causing silently transformed names:
	//     https://github.com/GoogleCloudPlatform/gcloud-golang/issues/111
	if !utf8.ValidString(attrs.Name) {
		return nil, errors.New("Invalid object name: not valid UTF-8")
	}

	// Create and initialize the wrapped writer.
	wrapped := storage.NewWriter(authContext, b.name, attrs.Name)
	wrapped.ObjectAttrs = *attrs

	w := &objectWriter{
		wrapped: wrapped,
	}

	return w, nil
}

func (b *bucket) DeleteObject(ctx context.Context, name string) error {
	authContext := cloud.WithContext(ctx, b.projID, b.client)
	return storage.DeleteObject(authContext, b.name, name)
}
<|endoftext|>


<commit_before>/*
Copyright (C) 2016 Red Hat, Inc.

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 cluster

import (
	"fmt"

	"github.com/jimmidyson/minishift/pkg/minikube/constants"
)

// Kill any running instances.
var stopCommand = "sudo killall openshift || true"

var startCommandFmtStr = `
# Run with nohup so it stays up. Redirect logs to useful places.
cd /var/lib/minishift;
if [ ! -f openshift.local.config/master/master-config.yaml ]; then
    sudo /usr/local/bin/openshift start --listen=https://0.0.0.0:%d --cors-allowed-origins=.* --master=https://%s:%d --write-config=openshift.local.config;
    sudo /usr/local/bin/openshift ex config patch openshift.local.config/master/master-config.yaml --patch='{"routingConfig": {"subdomain": "%s.xip.io"}}' > /tmp/master-config.yaml;
    sudo mv /tmp/master-config.yaml openshift.local.config/master/master-config.yaml;
    sudo /usr/local/bin/openshift ex config patch openshift.local.config/node-minishift/node-config.yaml --patch='{"nodeIP": "%s"}}' > /tmp/node-config.yaml;
    sudo mv /tmp/node-config.yaml openshift.local.config/node-minishift/node-config.yaml;
fi;
sudo sh -c 'PATH=/usr/local/sbin:$PATH nohup /usr/local/bin/openshift start --master-config=openshift.local.config/master/master-config.yaml --node-config=openshift.local.config/node-$(hostname | tr '[:upper:]' '[:lower:]')/node-config.yaml> %s 2> %s < /dev/null &'
until $(curl --output /dev/null --silent --fail -k https://localhost:%d/healthz/ready); do
    printf '.'
    sleep 1
done;
sudo /usr/local/bin/openshift admin policy add-cluster-role-to-user cluster-admin admin --config=openshift.local.config/master/admin.kubeconfig
`

var (
	logsCommand      = fmt.Sprintf("tail -n +1 %s %s", constants.RemoteOpenShiftErrPath, constants.RemoteOpenShiftOutPath)
	getCACertCommand = fmt.Sprintf("cat %s", constants.RemoteOpenShiftCAPath)
)

func GetStartCommand(ip string) string {
	return fmt.Sprintf(startCommandFmtStr, constants.APIServerPort, ip, constants.APIServerPort, ip, ip, constants.RemoteOpenShiftErrPath, constants.RemoteOpenShiftOutPath, constants.APIServerPort)
}
<commit_msg>Enable hostPath provisioner<commit_after>/*
Copyright (C) 2016 Red Hat, Inc.

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 cluster

import (
	"fmt"

	"github.com/jimmidyson/minishift/pkg/minikube/constants"
)

// Kill any running instances.
var stopCommand = "sudo killall openshift || true"

var startCommandFmtStr = `
# Run with nohup so it stays up. Redirect logs to useful places.
cd /var/lib/minishift;
if [ ! -f openshift.local.config/master/master-config.yaml ]; then
    sudo /usr/local/bin/openshift start --listen=https://0.0.0.0:%d --cors-allowed-origins=.* --master=https://%s:%d --write-config=openshift.local.config;
    sudo /usr/local/bin/openshift ex config patch openshift.local.config/master/master-config.yaml --patch='{"routingConfig": {"subdomain": "%s.xip.io"}}' > /tmp/master-config.yaml;
    sudo mv /tmp/master-config.yaml openshift.local.config/master/master-config.yaml;
		sudo /usr/local/bin/openshift ex config patch openshift.local.config/master/master-config.yaml --patch='{"kubernetesMasterConfig": {"controllerArguments": {"enable-hostpath-provisioner": ["true"]}}}' > /tmp/master-config.yaml;
    sudo mv /tmp/master-config.yaml openshift.local.config/master/master-config.yaml;
    sudo /usr/local/bin/openshift ex config patch openshift.local.config/node-minishift/node-config.yaml --patch='{"nodeIP": "%s"}}' > /tmp/node-config.yaml;
    sudo mv /tmp/node-config.yaml openshift.local.config/node-minishift/node-config.yaml;
fi;
sudo sh -c 'PATH=/usr/local/sbin:$PATH nohup /usr/local/bin/openshift start --master-config=openshift.local.config/master/master-config.yaml --node-config=openshift.local.config/node-$(hostname | tr '[:upper:]' '[:lower:]')/node-config.yaml> %s 2> %s < /dev/null &'
until $(curl --output /dev/null --silent --fail -k https://localhost:%d/healthz/ready); do
    printf '.'
    sleep 1
done;
sudo /usr/local/bin/openshift admin policy add-cluster-role-to-user cluster-admin admin --config=openshift.local.config/master/admin.kubeconfig
`

var (
	logsCommand      = fmt.Sprintf("tail -n +1 %s %s", constants.RemoteOpenShiftErrPath, constants.RemoteOpenShiftOutPath)
	getCACertCommand = fmt.Sprintf("cat %s", constants.RemoteOpenShiftCAPath)
)

func GetStartCommand(ip string) string {
	return fmt.Sprintf(startCommandFmtStr, constants.APIServerPort, ip, constants.APIServerPort, ip, ip, constants.RemoteOpenShiftErrPath, constants.RemoteOpenShiftOutPath, constants.APIServerPort)
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Rook 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 object

import (
	"fmt"
	"strings"

	"github.com/pkg/errors"
	cephv1 "github.com/rook/rook/pkg/apis/ceph.rook.io/v1"
	cephconfig "github.com/rook/rook/pkg/operator/ceph/config"
	"github.com/rook/rook/pkg/operator/ceph/controller"
	"github.com/rook/rook/pkg/operator/k8sutil"
	apps "k8s.io/api/apps/v1"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/intstr"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
)

const (
	livenessProbePath = "/swift/healthcheck"
)

func (c *clusterConfig) createDeployment(rgwConfig *rgwConfig) (*apps.Deployment, error) {
	pod, err := c.makeRGWPodSpec(rgwConfig)
	if err != nil {
		return nil, err
	}
	replicas := int32(1)
	d := &apps.Deployment{
		ObjectMeta: metav1.ObjectMeta{
			Name:      rgwConfig.ResourceName,
			Namespace: c.store.Namespace,
			Labels:    getLabels(c.store.Name, c.store.Namespace, true),
		},
		Spec: apps.DeploymentSpec{
			Selector: &metav1.LabelSelector{
				MatchLabels: getLabels(c.store.Name, c.store.Namespace, false),
			},
			Template: pod,
			Replicas: &replicas,
			Strategy: apps.DeploymentStrategy{
				Type: apps.RecreateDeploymentStrategyType,
			},
		},
	}
	k8sutil.AddRookVersionLabelToDeployment(d)
	c.store.Spec.Gateway.Annotations.ApplyToObjectMeta(&d.ObjectMeta)
	c.store.Spec.Gateway.Labels.ApplyToObjectMeta(&d.ObjectMeta)
	controller.AddCephVersionLabelToDeployment(c.clusterInfo.CephVersion, d)

	return d, nil
}

func (c *clusterConfig) makeRGWPodSpec(rgwConfig *rgwConfig) (v1.PodTemplateSpec, error) {
	podSpec := v1.PodSpec{
		InitContainers: []v1.Container{
			c.makeChownInitContainer(rgwConfig),
		},
		Containers:    []v1.Container{c.makeDaemonContainer(rgwConfig)},
		RestartPolicy: v1.RestartPolicyAlways,
		Volumes: append(
			controller.DaemonVolumes(c.DataPathMap, rgwConfig.ResourceName),
			c.mimeTypesVolume(),
		),
		HostNetwork:       c.clusterSpec.Network.IsHost(),
		PriorityClassName: c.store.Spec.Gateway.PriorityClassName,
	}

	// If the log collector is enabled we add the side-car container
	if c.clusterSpec.LogCollector.Enabled {
		shareProcessNamespace := true
		podSpec.ShareProcessNamespace = &shareProcessNamespace
		podSpec.Containers = append(podSpec.Containers, *controller.LogCollectorContainer(strings.TrimPrefix(generateCephXUser(fmt.Sprintf("ceph-client.%s", rgwConfig.ResourceName)), "client."), c.clusterInfo.Namespace, *c.clusterSpec))
	}

	// Replace default unreachable node toleration
	k8sutil.AddUnreachableNodeToleration(&podSpec)

	// Set the ssl cert if specified
	if c.store.Spec.Gateway.SSLCertificateRef != "" {
		// Keep the SSL secret as secure as possible in the container. Give only user read perms.
		// Because the Secret mount is owned by "root" and fsGroup breaks on OCP since we cannot predict it
		// Also, we don't want to change the SCC for fsGroup to RunAsAny since it has a major broader impact
		// Let's open the permissions a bit more so that everyone can read the cert.
		userReadOnly := int32(0444)
		certVol := v1.Volume{
			Name: certVolumeName,
			VolumeSource: v1.VolumeSource{
				Secret: &v1.SecretVolumeSource{
					SecretName: c.store.Spec.Gateway.SSLCertificateRef,
					Items: []v1.KeyToPath{
						{Key: certKeyName, Path: certFilename, Mode: &userReadOnly},
					}}}}
		podSpec.Volumes = append(podSpec.Volumes, certVol)
	}

	// If host networking is not enabled, preferred pod anti-affinity is added to the rgw daemons
	labels := getLabels(c.store.Name, c.store.Namespace, false)
	k8sutil.SetNodeAntiAffinityForPod(&podSpec, c.store.Spec.Gateway.Placement, c.clusterSpec.Network.IsHost(), labels, nil)

	podTemplateSpec := v1.PodTemplateSpec{
		ObjectMeta: metav1.ObjectMeta{
			Name:   rgwConfig.ResourceName,
			Labels: getLabels(c.store.Name, c.store.Namespace, true),
		},
		Spec: podSpec,
	}
	c.store.Spec.Gateway.Annotations.ApplyToObjectMeta(&podTemplateSpec.ObjectMeta)
	c.store.Spec.Gateway.Labels.ApplyToObjectMeta(&podTemplateSpec.ObjectMeta)

	if c.clusterSpec.Network.IsHost() {
		podTemplateSpec.Spec.DNSPolicy = v1.DNSClusterFirstWithHostNet
	} else if c.clusterSpec.Network.IsMultus() {
		if err := k8sutil.ApplyMultus(c.clusterSpec.Network.NetworkSpec, &podTemplateSpec.ObjectMeta); err != nil {
			return podTemplateSpec, err
		}
	}

	return podTemplateSpec, nil
}

func (c *clusterConfig) makeChownInitContainer(rgwConfig *rgwConfig) v1.Container {
	return controller.ChownCephDataDirsInitContainer(
		*c.DataPathMap,
		c.clusterSpec.CephVersion.Image,
		controller.DaemonVolumeMounts(c.DataPathMap, rgwConfig.ResourceName),
		c.store.Spec.Gateway.Resources,
		controller.PodSecurityContext(),
	)
}

func (c *clusterConfig) makeDaemonContainer(rgwConfig *rgwConfig) v1.Container {
	// start the rgw daemon in the foreground
	container := v1.Container{
		Name:  "rgw",
		Image: c.clusterSpec.CephVersion.Image,
		Command: []string{
			"radosgw",
		},
		Args: append(
			controller.DaemonFlags(c.clusterInfo, c.clusterSpec,
				strings.TrimPrefix(generateCephXUser(rgwConfig.ResourceName), "client.")),
			"--foreground",
			cephconfig.NewFlag("rgw frontends", fmt.Sprintf("%s %s", rgwFrontendName, c.portString())),
			cephconfig.NewFlag("host", controller.ContainerEnvVarReference(k8sutil.PodNameEnvVar)),
			cephconfig.NewFlag("rgw-mime-types-file", mimeTypesMountPath()),
			cephconfig.NewFlag("rgw realm", rgwConfig.Realm),
			cephconfig.NewFlag("rgw zonegroup", rgwConfig.ZoneGroup),
			cephconfig.NewFlag("rgw zone", rgwConfig.Zone),
		),
		VolumeMounts: append(
			controller.DaemonVolumeMounts(c.DataPathMap, rgwConfig.ResourceName),
			c.mimeTypesVolumeMount(),
		),
		Env:             controller.DaemonEnvVars(c.clusterSpec.CephVersion.Image),
		Resources:       c.store.Spec.Gateway.Resources,
		LivenessProbe:   c.generateLiveProbe(),
		SecurityContext: controller.PodSecurityContext(),
	}

	// If the liveness probe is enabled
	configureLivenessProbe(&container, c.store.Spec.HealthCheck)
	if c.store.Spec.Gateway.SSLCertificateRef != "" {
		// Add a volume mount for the ssl certificate
		mount := v1.VolumeMount{Name: certVolumeName, MountPath: certDir, ReadOnly: true}
		container.VolumeMounts = append(container.VolumeMounts, mount)
	}

	return container
}

// configureLivenessProbe returns the desired liveness probe for a given daemon
func configureLivenessProbe(container *v1.Container, healthCheck cephv1.BucketHealthCheckSpec) {
	if ok := healthCheck.LivenessProbe; ok != nil {
		if !healthCheck.LivenessProbe.Disabled {
			probe := healthCheck.LivenessProbe.Probe
			// If the spec value is empty, let's use a default
			if probe != nil {
				container.LivenessProbe = probe
			}
		} else {
			container.LivenessProbe = nil
		}
	}
}

func (c *clusterConfig) generateLiveProbe() *v1.Probe {
	return &v1.Probe{
		Handler: v1.Handler{
			HTTPGet: &v1.HTTPGetAction{
				Path:   livenessProbePath,
				Port:   c.generateLiveProbePort(),
				Scheme: c.generateLiveProbeScheme(),
			},
		},
		InitialDelaySeconds: 10,
	}
}

func (c *clusterConfig) generateLiveProbeScheme() v1.URIScheme {
	// Default to HTTP
	uriScheme := v1.URISchemeHTTP

	// If rgw is configured to use a secured port we need get on https://
	// Only do this when the Non-SSL port is not used
	if c.store.Spec.Gateway.Port == 0 && c.store.Spec.Gateway.SecurePort != 0 && c.store.Spec.Gateway.SSLCertificateRef != "" {
		uriScheme = v1.URISchemeHTTPS
	}

	return uriScheme
}

func (c *clusterConfig) generateLiveProbePort() intstr.IntOrString {
	// The port the liveness probe needs to probe
	// Assume we run on SDN by default
	port := intstr.FromInt(int(rgwPortInternalPort))

	// If Host Networking is enabled, the port from the spec must be reflected
	if c.clusterSpec.Network.IsHost() {
		if c.store.Spec.Gateway.Port == 0 && c.store.Spec.Gateway.SecurePort != 0 && c.store.Spec.Gateway.SSLCertificateRef != "" {
			port = intstr.FromInt(int(c.store.Spec.Gateway.SecurePort))
		} else {
			port = intstr.FromInt(int(c.store.Spec.Gateway.Port))
		}
	}

	return port
}

func (c *clusterConfig) generateService(cephObjectStore *cephv1.CephObjectStore) *v1.Service {
	labels := getLabels(cephObjectStore.Name, cephObjectStore.Namespace, true)
	svc := &v1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:      instanceName(cephObjectStore.Name),
			Namespace: cephObjectStore.Namespace,
			Labels:    labels,
		},
	}

	if c.clusterSpec.Network.IsHost() {
		svc.Spec.ClusterIP = v1.ClusterIPNone
	}

	destPort := c.generateLiveProbePort()

	// When the cluster is external we must use the same one as the gateways are listening on
	if c.clusterSpec.External.Enable {
		destPort.IntVal = cephObjectStore.Spec.Gateway.Port
	} else {
		// If the cluster is not external we add the Selector
		svc.Spec = v1.ServiceSpec{
			Selector: labels,
		}
	}
	addPort(svc, "http", cephObjectStore.Spec.Gateway.Port, destPort.IntVal)
	addPort(svc, "https", cephObjectStore.Spec.Gateway.SecurePort, cephObjectStore.Spec.Gateway.SecurePort)

	return svc
}

func (c *clusterConfig) generateEndpoint(cephObjectStore *cephv1.CephObjectStore) *v1.Endpoints {
	labels := getLabels(cephObjectStore.Name, cephObjectStore.Namespace, true)

	endpoints := &v1.Endpoints{
		ObjectMeta: metav1.ObjectMeta{
			Name:      instanceName(cephObjectStore.Name),
			Namespace: cephObjectStore.Namespace,
			Labels:    labels,
		},
		Subsets: []v1.EndpointSubset{
			{
				Addresses: cephObjectStore.Spec.Gateway.ExternalRgwEndpoints,
			},
		},
	}

	addPortToEndpoint(endpoints, "http", cephObjectStore.Spec.Gateway.Port)
	addPortToEndpoint(endpoints, "https", cephObjectStore.Spec.Gateway.SecurePort)

	return endpoints
}

func (c *clusterConfig) reconcileExternalEndpoint(cephObjectStore *cephv1.CephObjectStore) error {
	logger.Info("reconciling external object store service")

	endpoint := c.generateEndpoint(cephObjectStore)
	// Set owner ref to the parent object
	err := controllerutil.SetControllerReference(cephObjectStore, endpoint, c.scheme)
	if err != nil {
		return errors.Wrap(err, "failed to set owner reference to ceph object store endpoint")
	}

	_, err = k8sutil.CreateOrUpdateEndpoint(c.context.Clientset, cephObjectStore.Namespace, endpoint)
	if err != nil {
		return errors.Wrapf(err, "failed to create or update object store %q endpoint", cephObjectStore.Name)
	}

	return nil
}

func (c *clusterConfig) reconcileService(cephObjectStore *cephv1.CephObjectStore) (string, error) {
	service := c.generateService(cephObjectStore)
	// Set owner ref to the parent object
	err := controllerutil.SetControllerReference(cephObjectStore, service, c.scheme)
	if err != nil {
		return "", errors.Wrap(err, "failed to set owner reference to ceph object store service")
	}

	svc, err := k8sutil.CreateOrUpdateService(c.context.Clientset, cephObjectStore.Namespace, service)
	if err != nil {
		return "", errors.Wrapf(err, "failed to create or update object store %q service", cephObjectStore.Name)
	}

	logger.Infof("ceph object store gateway service running at %s", svc.Spec.ClusterIP)

	return svc.Spec.ClusterIP, nil
}

func addPort(service *v1.Service, name string, port, destPort int32) {
	if port == 0 || destPort == 0 {
		return
	}
	service.Spec.Ports = append(service.Spec.Ports, v1.ServicePort{
		Name:       name,
		Port:       port,
		TargetPort: intstr.FromInt(int(destPort)),
		Protocol:   v1.ProtocolTCP,
	})
}

func addPortToEndpoint(endpoints *v1.Endpoints, name string, port int32) {
	if port == 0 {
		return
	}
	endpoints.Subsets[0].Ports = append(endpoints.Subsets[0].Ports, v1.EndpointPort{
		Name:     name,
		Port:     port,
		Protocol: v1.ProtocolTCP,
	},
	)
}

func getLabels(name, namespace string, includeNewLabels bool) map[string]string {
	labels := controller.CephDaemonAppLabels(AppName, namespace, "rgw", name, includeNewLabels)
	labels["rook_object_store"] = name
	return labels
}
<commit_msg>ceph: rgw service selector should not change<commit_after>/*
Copyright 2018 The Rook 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 object

import (
	"fmt"
	"strings"

	"github.com/pkg/errors"
	cephv1 "github.com/rook/rook/pkg/apis/ceph.rook.io/v1"
	cephconfig "github.com/rook/rook/pkg/operator/ceph/config"
	"github.com/rook/rook/pkg/operator/ceph/controller"
	"github.com/rook/rook/pkg/operator/k8sutil"
	apps "k8s.io/api/apps/v1"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/intstr"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
)

const (
	livenessProbePath = "/swift/healthcheck"
)

func (c *clusterConfig) createDeployment(rgwConfig *rgwConfig) (*apps.Deployment, error) {
	pod, err := c.makeRGWPodSpec(rgwConfig)
	if err != nil {
		return nil, err
	}
	replicas := int32(1)
	d := &apps.Deployment{
		ObjectMeta: metav1.ObjectMeta{
			Name:      rgwConfig.ResourceName,
			Namespace: c.store.Namespace,
			Labels:    getLabels(c.store.Name, c.store.Namespace, true),
		},
		Spec: apps.DeploymentSpec{
			Selector: &metav1.LabelSelector{
				MatchLabels: getLabels(c.store.Name, c.store.Namespace, false),
			},
			Template: pod,
			Replicas: &replicas,
			Strategy: apps.DeploymentStrategy{
				Type: apps.RecreateDeploymentStrategyType,
			},
		},
	}
	k8sutil.AddRookVersionLabelToDeployment(d)
	c.store.Spec.Gateway.Annotations.ApplyToObjectMeta(&d.ObjectMeta)
	c.store.Spec.Gateway.Labels.ApplyToObjectMeta(&d.ObjectMeta)
	controller.AddCephVersionLabelToDeployment(c.clusterInfo.CephVersion, d)

	return d, nil
}

func (c *clusterConfig) makeRGWPodSpec(rgwConfig *rgwConfig) (v1.PodTemplateSpec, error) {
	podSpec := v1.PodSpec{
		InitContainers: []v1.Container{
			c.makeChownInitContainer(rgwConfig),
		},
		Containers:    []v1.Container{c.makeDaemonContainer(rgwConfig)},
		RestartPolicy: v1.RestartPolicyAlways,
		Volumes: append(
			controller.DaemonVolumes(c.DataPathMap, rgwConfig.ResourceName),
			c.mimeTypesVolume(),
		),
		HostNetwork:       c.clusterSpec.Network.IsHost(),
		PriorityClassName: c.store.Spec.Gateway.PriorityClassName,
	}

	// If the log collector is enabled we add the side-car container
	if c.clusterSpec.LogCollector.Enabled {
		shareProcessNamespace := true
		podSpec.ShareProcessNamespace = &shareProcessNamespace
		podSpec.Containers = append(podSpec.Containers, *controller.LogCollectorContainer(strings.TrimPrefix(generateCephXUser(fmt.Sprintf("ceph-client.%s", rgwConfig.ResourceName)), "client."), c.clusterInfo.Namespace, *c.clusterSpec))
	}

	// Replace default unreachable node toleration
	k8sutil.AddUnreachableNodeToleration(&podSpec)

	// Set the ssl cert if specified
	if c.store.Spec.Gateway.SSLCertificateRef != "" {
		// Keep the SSL secret as secure as possible in the container. Give only user read perms.
		// Because the Secret mount is owned by "root" and fsGroup breaks on OCP since we cannot predict it
		// Also, we don't want to change the SCC for fsGroup to RunAsAny since it has a major broader impact
		// Let's open the permissions a bit more so that everyone can read the cert.
		userReadOnly := int32(0444)
		certVol := v1.Volume{
			Name: certVolumeName,
			VolumeSource: v1.VolumeSource{
				Secret: &v1.SecretVolumeSource{
					SecretName: c.store.Spec.Gateway.SSLCertificateRef,
					Items: []v1.KeyToPath{
						{Key: certKeyName, Path: certFilename, Mode: &userReadOnly},
					}}}}
		podSpec.Volumes = append(podSpec.Volumes, certVol)
	}

	// If host networking is not enabled, preferred pod anti-affinity is added to the rgw daemons
	labels := getLabels(c.store.Name, c.store.Namespace, false)
	k8sutil.SetNodeAntiAffinityForPod(&podSpec, c.store.Spec.Gateway.Placement, c.clusterSpec.Network.IsHost(), labels, nil)

	podTemplateSpec := v1.PodTemplateSpec{
		ObjectMeta: metav1.ObjectMeta{
			Name:   rgwConfig.ResourceName,
			Labels: getLabels(c.store.Name, c.store.Namespace, true),
		},
		Spec: podSpec,
	}
	c.store.Spec.Gateway.Annotations.ApplyToObjectMeta(&podTemplateSpec.ObjectMeta)
	c.store.Spec.Gateway.Labels.ApplyToObjectMeta(&podTemplateSpec.ObjectMeta)

	if c.clusterSpec.Network.IsHost() {
		podTemplateSpec.Spec.DNSPolicy = v1.DNSClusterFirstWithHostNet
	} else if c.clusterSpec.Network.IsMultus() {
		if err := k8sutil.ApplyMultus(c.clusterSpec.Network.NetworkSpec, &podTemplateSpec.ObjectMeta); err != nil {
			return podTemplateSpec, err
		}
	}

	return podTemplateSpec, nil
}

func (c *clusterConfig) makeChownInitContainer(rgwConfig *rgwConfig) v1.Container {
	return controller.ChownCephDataDirsInitContainer(
		*c.DataPathMap,
		c.clusterSpec.CephVersion.Image,
		controller.DaemonVolumeMounts(c.DataPathMap, rgwConfig.ResourceName),
		c.store.Spec.Gateway.Resources,
		controller.PodSecurityContext(),
	)
}

func (c *clusterConfig) makeDaemonContainer(rgwConfig *rgwConfig) v1.Container {
	// start the rgw daemon in the foreground
	container := v1.Container{
		Name:  "rgw",
		Image: c.clusterSpec.CephVersion.Image,
		Command: []string{
			"radosgw",
		},
		Args: append(
			controller.DaemonFlags(c.clusterInfo, c.clusterSpec,
				strings.TrimPrefix(generateCephXUser(rgwConfig.ResourceName), "client.")),
			"--foreground",
			cephconfig.NewFlag("rgw frontends", fmt.Sprintf("%s %s", rgwFrontendName, c.portString())),
			cephconfig.NewFlag("host", controller.ContainerEnvVarReference(k8sutil.PodNameEnvVar)),
			cephconfig.NewFlag("rgw-mime-types-file", mimeTypesMountPath()),
			cephconfig.NewFlag("rgw realm", rgwConfig.Realm),
			cephconfig.NewFlag("rgw zonegroup", rgwConfig.ZoneGroup),
			cephconfig.NewFlag("rgw zone", rgwConfig.Zone),
		),
		VolumeMounts: append(
			controller.DaemonVolumeMounts(c.DataPathMap, rgwConfig.ResourceName),
			c.mimeTypesVolumeMount(),
		),
		Env:             controller.DaemonEnvVars(c.clusterSpec.CephVersion.Image),
		Resources:       c.store.Spec.Gateway.Resources,
		LivenessProbe:   c.generateLiveProbe(),
		SecurityContext: controller.PodSecurityContext(),
	}

	// If the liveness probe is enabled
	configureLivenessProbe(&container, c.store.Spec.HealthCheck)
	if c.store.Spec.Gateway.SSLCertificateRef != "" {
		// Add a volume mount for the ssl certificate
		mount := v1.VolumeMount{Name: certVolumeName, MountPath: certDir, ReadOnly: true}
		container.VolumeMounts = append(container.VolumeMounts, mount)
	}

	return container
}

// configureLivenessProbe returns the desired liveness probe for a given daemon
func configureLivenessProbe(container *v1.Container, healthCheck cephv1.BucketHealthCheckSpec) {
	if ok := healthCheck.LivenessProbe; ok != nil {
		if !healthCheck.LivenessProbe.Disabled {
			probe := healthCheck.LivenessProbe.Probe
			// If the spec value is empty, let's use a default
			if probe != nil {
				container.LivenessProbe = probe
			}
		} else {
			container.LivenessProbe = nil
		}
	}
}

func (c *clusterConfig) generateLiveProbe() *v1.Probe {
	return &v1.Probe{
		Handler: v1.Handler{
			HTTPGet: &v1.HTTPGetAction{
				Path:   livenessProbePath,
				Port:   c.generateLiveProbePort(),
				Scheme: c.generateLiveProbeScheme(),
			},
		},
		InitialDelaySeconds: 10,
	}
}

func (c *clusterConfig) generateLiveProbeScheme() v1.URIScheme {
	// Default to HTTP
	uriScheme := v1.URISchemeHTTP

	// If rgw is configured to use a secured port we need get on https://
	// Only do this when the Non-SSL port is not used
	if c.store.Spec.Gateway.Port == 0 && c.store.Spec.Gateway.SecurePort != 0 && c.store.Spec.Gateway.SSLCertificateRef != "" {
		uriScheme = v1.URISchemeHTTPS
	}

	return uriScheme
}

func (c *clusterConfig) generateLiveProbePort() intstr.IntOrString {
	// The port the liveness probe needs to probe
	// Assume we run on SDN by default
	port := intstr.FromInt(int(rgwPortInternalPort))

	// If Host Networking is enabled, the port from the spec must be reflected
	if c.clusterSpec.Network.IsHost() {
		if c.store.Spec.Gateway.Port == 0 && c.store.Spec.Gateway.SecurePort != 0 && c.store.Spec.Gateway.SSLCertificateRef != "" {
			port = intstr.FromInt(int(c.store.Spec.Gateway.SecurePort))
		} else {
			port = intstr.FromInt(int(c.store.Spec.Gateway.Port))
		}
	}

	return port
}

func (c *clusterConfig) generateService(cephObjectStore *cephv1.CephObjectStore) *v1.Service {
	svc := &v1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:      instanceName(cephObjectStore.Name),
			Namespace: cephObjectStore.Namespace,
			Labels:    getLabels(cephObjectStore.Name, cephObjectStore.Namespace, true),
		},
	}

	if c.clusterSpec.Network.IsHost() {
		svc.Spec.ClusterIP = v1.ClusterIPNone
	}

	destPort := c.generateLiveProbePort()

	// When the cluster is external we must use the same one as the gateways are listening on
	if c.clusterSpec.External.Enable {
		destPort.IntVal = cephObjectStore.Spec.Gateway.Port
	} else {
		// If the cluster is not external we add the Selector
		svc.Spec = v1.ServiceSpec{
			Selector: getLabels(cephObjectStore.Name, cephObjectStore.Namespace, false),
		}
	}
	addPort(svc, "http", cephObjectStore.Spec.Gateway.Port, destPort.IntVal)
	addPort(svc, "https", cephObjectStore.Spec.Gateway.SecurePort, cephObjectStore.Spec.Gateway.SecurePort)

	return svc
}

func (c *clusterConfig) generateEndpoint(cephObjectStore *cephv1.CephObjectStore) *v1.Endpoints {
	labels := getLabels(cephObjectStore.Name, cephObjectStore.Namespace, true)

	endpoints := &v1.Endpoints{
		ObjectMeta: metav1.ObjectMeta{
			Name:      instanceName(cephObjectStore.Name),
			Namespace: cephObjectStore.Namespace,
			Labels:    labels,
		},
		Subsets: []v1.EndpointSubset{
			{
				Addresses: cephObjectStore.Spec.Gateway.ExternalRgwEndpoints,
			},
		},
	}

	addPortToEndpoint(endpoints, "http", cephObjectStore.Spec.Gateway.Port)
	addPortToEndpoint(endpoints, "https", cephObjectStore.Spec.Gateway.SecurePort)

	return endpoints
}

func (c *clusterConfig) reconcileExternalEndpoint(cephObjectStore *cephv1.CephObjectStore) error {
	logger.Info("reconciling external object store service")

	endpoint := c.generateEndpoint(cephObjectStore)
	// Set owner ref to the parent object
	err := controllerutil.SetControllerReference(cephObjectStore, endpoint, c.scheme)
	if err != nil {
		return errors.Wrap(err, "failed to set owner reference to ceph object store endpoint")
	}

	_, err = k8sutil.CreateOrUpdateEndpoint(c.context.Clientset, cephObjectStore.Namespace, endpoint)
	if err != nil {
		return errors.Wrapf(err, "failed to create or update object store %q endpoint", cephObjectStore.Name)
	}

	return nil
}

func (c *clusterConfig) reconcileService(cephObjectStore *cephv1.CephObjectStore) (string, error) {
	service := c.generateService(cephObjectStore)
	// Set owner ref to the parent object
	err := controllerutil.SetControllerReference(cephObjectStore, service, c.scheme)
	if err != nil {
		return "", errors.Wrap(err, "failed to set owner reference to ceph object store service")
	}

	svc, err := k8sutil.CreateOrUpdateService(c.context.Clientset, cephObjectStore.Namespace, service)
	if err != nil {
		return "", errors.Wrapf(err, "failed to create or update object store %q service", cephObjectStore.Name)
	}

	logger.Infof("ceph object store gateway service running at %s", svc.Spec.ClusterIP)

	return svc.Spec.ClusterIP, nil
}

func addPort(service *v1.Service, name string, port, destPort int32) {
	if port == 0 || destPort == 0 {
		return
	}
	service.Spec.Ports = append(service.Spec.Ports, v1.ServicePort{
		Name:       name,
		Port:       port,
		TargetPort: intstr.FromInt(int(destPort)),
		Protocol:   v1.ProtocolTCP,
	})
}

func addPortToEndpoint(endpoints *v1.Endpoints, name string, port int32) {
	if port == 0 {
		return
	}
	endpoints.Subsets[0].Ports = append(endpoints.Subsets[0].Ports, v1.EndpointPort{
		Name:     name,
		Port:     port,
		Protocol: v1.ProtocolTCP,
	},
	)
}

func getLabels(name, namespace string, includeNewLabels bool) map[string]string {
	labels := controller.CephDaemonAppLabels(AppName, namespace, "rgw", name, includeNewLabels)
	labels["rook_object_store"] = name
	return labels
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Rook 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 to manage a Minio object store.
package minio

import (
	"fmt"
	"reflect"

	"github.com/coreos/pkg/capnslog"
	opkit "github.com/rook/operator-kit"
	miniov1alpha1 "github.com/rook/rook/pkg/apis/minio.rook.io/v1alpha1"
	"github.com/rook/rook/pkg/clusterd"
	"github.com/rook/rook/pkg/operator/k8sutil"
	"k8s.io/api/apps/v1beta2"
	"k8s.io/api/core/v1"
	apiextensionsv1beta1 "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1beta1"
	"k8s.io/apimachinery/pkg/api/resource"
	meta_v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/tools/cache"
)

// TODO: A lot of these constants are specific to the KubeCon demo. Let's
// revisit these and determine what should be specified in the resource spec.
const (
	customResourceName       = "objectstore"
	customResourceNamePlural = "objectstores"
	minioCtrName             = "minio"
	minioLabel               = "minio"
	minioPVCName             = "minio-pvc"
	minioServerSuffixFmt     = "%s.svc.cluster.local"
	minioVolumeName          = "data"
	objectStoreDataDir       = "/data"
)

var logger = capnslog.NewPackageLogger("github.com/rook/rook", "minio-op-object")

// ObjectStoreResource represents the object store custom resource
var ObjectStoreResource = opkit.CustomResource{
	Name:    customResourceName,
	Plural:  customResourceNamePlural,
	Group:   miniov1alpha1.CustomResourceGroup,
	Version: miniov1alpha1.Version,
	Scope:   apiextensionsv1beta1.NamespaceScoped,
	Kind:    reflect.TypeOf(miniov1alpha1.ObjectStore{}).Name(),
}

// MinioController represents a controller object for object store custom resources
type MinioController struct {
	context   *clusterd.Context
	rookImage string
}

// NewMinioController create controller for watching object store custom resources created
func NewMinioController(context *clusterd.Context, rookImage string) *MinioController {
	return &MinioController{
		context:   context,
		rookImage: rookImage,
	}
}

// StartWatch watches for instances of ObjectStore custom resources and acts on them
func (c *MinioController) StartWatch(namespace string, stopCh chan struct{}) error {
	resourceHandlerFuncs := cache.ResourceEventHandlerFuncs{
		AddFunc:    c.onAdd,
		UpdateFunc: c.onUpdate,
		DeleteFunc: c.onDelete,
	}

	logger.Infof("start watching object store resources in namespace %s", namespace)
	watcher := opkit.NewWatcher(ObjectStoreResource, namespace, resourceHandlerFuncs, c.context.RookClientset.MinioV1alpha1().RESTClient())
	go watcher.Watch(&miniov1alpha1.ObjectStore{}, stopCh)

	return nil
}

func (c *MinioController) makeMinioHeadlessService(name, namespace string, spec miniov1alpha1.ObjectStoreSpec, ownerRef meta_v1.OwnerReference) (*v1.Service, error) {
	coreV1Client := c.context.Clientset.CoreV1()

	svc, err := coreV1Client.Services(namespace).Create(&v1.Service{
		ObjectMeta: meta_v1.ObjectMeta{
			Name:            name,
			Namespace:       namespace,
			Labels:          map[string]string{k8sutil.AppAttr: minioLabel},
			OwnerReferences: []meta_v1.OwnerReference{ownerRef},
		},
		Spec: v1.ServiceSpec{
			Selector:  map[string]string{k8sutil.AppAttr: minioLabel},
			Ports:     []v1.ServicePort{{Port: spec.Port}},
			ClusterIP: v1.ClusterIPNone,
		},
	})

	return svc, err
}

func (c *MinioController) buildMinioCtrArgs(statefulSetPrefix, headlessServiceName, namespace string, serverCount int32) []string {
	args := []string{"server"}
	for i := int32(0); i < serverCount; i++ {
		serverAddress := fmt.Sprintf("http://%s-%d.%s.%s%s", statefulSetPrefix, i, headlessServiceName, fmt.Sprintf(minioServerSuffixFmt, namespace), objectStoreDataDir)
		args = append(args, serverAddress)
	}

	logger.Infof("Building Minio container args: %v", args)
	return args
}

func (c *MinioController) makeMinioPodSpec(name, namespace string, ctrName string, ctrImage string, port int32, envVars map[string]string, numServers int32) v1.PodTemplateSpec {
	var env []v1.EnvVar
	for k, v := range envVars {
		env = append(env, v1.EnvVar{Name: k, Value: v})
	}

	podSpec := v1.PodTemplateSpec{
		ObjectMeta: meta_v1.ObjectMeta{
			Name:      name,
			Namespace: namespace,
			Labels:    map[string]string{k8sutil.AppAttr: minioLabel},
		},
		Spec: v1.PodSpec{
			Containers: []v1.Container{
				{
					Name:    ctrName,
					Image:   ctrImage,
					Env:     env,
					Command: []string{"/usr/bin/minio"},
					Ports:   []v1.ContainerPort{{ContainerPort: port}},
					Args:    c.buildMinioCtrArgs(name, name, namespace, numServers),
					VolumeMounts: []v1.VolumeMount{
						{
							Name:      minioVolumeName,
							MountPath: objectStoreDataDir,
						},
					},
				},
			},
			Volumes: []v1.Volume{
				{
					Name: minioVolumeName,
					VolumeSource: v1.VolumeSource{
						PersistentVolumeClaim: &v1.PersistentVolumeClaimVolumeSource{
							ClaimName: minioPVCName,
						},
					},
				},
			},
		},
	}

	return podSpec
}

func (c *MinioController) getAccessCredentials(secretName, namespace string) (string, string, error) {
	coreV1Client := c.context.Clientset.CoreV1()
	var getOpts meta_v1.GetOptions
	val, err := coreV1Client.Secrets(namespace).Get(secretName, getOpts)
	if err != nil {
		logger.Errorf("Unable to get secret with name=%s in namespace=%s: %v", secretName, namespace, err)
		return "", "", err
	}

	return string(val.Data["username"]), string(val.Data["password"]), nil
}

func validateObjectStoreSpec(spec miniov1alpha1.ObjectStoreSpec) error {
	// Verify node count.
	count := spec.Storage.NodeCount
	if count < 4 || count%2 != 0 {
		return fmt.Errorf("Node count must be greater than 3 and even.")
	}

	// Verify sane port.
	if spec.Port < 1024 {
		return fmt.Errorf("Invalid port %d", spec.Port)
	}

	return nil
}

func (c *MinioController) makeMinioStatefulSet(name, namespace string, spec miniov1alpha1.ObjectStoreSpec, ownerRef meta_v1.OwnerReference) (*v1beta2.StatefulSet, error) {
	appsClient := c.context.Clientset.AppsV1beta2()

	accessKey, secretKey, err := c.getAccessCredentials(spec.Credentials.Name, spec.Credentials.Namespace)
	if err != nil {
		return nil, err
	}

	envVars := map[string]string{
		"MINIO_ACCESS_KEY": accessKey,
		"MINIO_SECRET_KEY": secretKey,
	}

	podSpec := c.makeMinioPodSpec(name, namespace, minioCtrName, c.rookImage, spec.Port, envVars, int32(spec.Storage.NodeCount))

	nodeCount := int32(spec.Storage.NodeCount)
	ss := v1beta2.StatefulSet{
		ObjectMeta: meta_v1.ObjectMeta{
			Name:            name,
			Namespace:       namespace,
			Labels:          map[string]string{k8sutil.AppAttr: minioLabel},
			OwnerReferences: []meta_v1.OwnerReference{ownerRef},
		},
		Spec: v1beta2.StatefulSetSpec{
			Replicas: &nodeCount,
			Selector: &meta_v1.LabelSelector{
				MatchLabels: map[string]string{k8sutil.AppAttr: minioLabel},
			},
			Template: podSpec,
			VolumeClaimTemplates: []v1.PersistentVolumeClaim{
				{
					ObjectMeta: meta_v1.ObjectMeta{
						Name:      minioVolumeName,
						Namespace: namespace,
					},
					Spec: v1.PersistentVolumeClaimSpec{
						AccessModes: []v1.PersistentVolumeAccessMode{v1.ReadWriteOnce},
						Resources: v1.ResourceRequirements{
							Requests: v1.ResourceList{
								v1.ResourceStorage: resource.MustParse(spec.StorageSize),
							},
						},
					},
				},
			},
			ServiceName: name,
			// TODO: liveness probe
		},
	}

	return appsClient.StatefulSets(namespace).Create(&ss)
}

func (c *MinioController) onAdd(obj interface{}) {
	objectstore := obj.(*miniov1alpha1.ObjectStore).DeepCopy()

	ownerRef := meta_v1.OwnerReference{
		APIVersion: ObjectStoreResource.Version,
		Kind:       ObjectStoreResource.Kind,
		Name:       objectstore.Namespace,
		UID:        types.UID(objectstore.ObjectMeta.UID),
	}

	// Validate object store config.
	err := validateObjectStoreSpec(objectstore.Spec)
	if err != nil {
		logger.Errorf("failed to validate object store config")
		return
	}

	// Create the headless service.
	logger.Infof("Creating Minio headless service %s in namespace %s.", objectstore.Name, objectstore.Namespace)
	_, err = c.makeMinioHeadlessService(objectstore.Name, objectstore.Namespace, objectstore.Spec, ownerRef)
	if err != nil {
		logger.Errorf("failed to create minio headless service: %v", err)
		return
	}
	logger.Infof("Finished creating Minio headless service %s in namespace %s.", objectstore.Name, objectstore.Namespace)

	// Create the stateful set.
	logger.Infof("Creating Minio stateful set %s.", objectstore.Name)
	_, err = c.makeMinioStatefulSet(objectstore.Name, objectstore.Namespace, objectstore.Spec, ownerRef)
	if err != nil {
		logger.Errorf("failed to create minio stateful set: %v", err)
		return
	}
	logger.Infof("Finished creating Minio stateful set %s in namespace %s.", objectstore.Name, objectstore.Namespace)
}

func (c *MinioController) onUpdate(oldObj, newObj interface{}) {
	oldStore := oldObj.(*miniov1alpha1.ObjectStore).DeepCopy()
	newStore := newObj.(*miniov1alpha1.ObjectStore).DeepCopy()

	_ = oldStore
	_ = newStore

	logger.Infof("Received update on object store %s in namespace %s. This is currently unsupported.", oldStore.Name, oldStore.Namespace)
}

func (c *MinioController) onDelete(obj interface{}) {
	objectstore := obj.(*miniov1alpha1.ObjectStore).DeepCopy()
	logger.Infof("Delete Minio object store %s", objectstore.Name)

	// Cleanup is handled by the owner references set in 'onAdd' and the k8s garbage collector.
}
<commit_msg>remove svc.cluster.local from serverAddress<commit_after>/*
Copyright 2018 The Rook 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 to manage a Minio object store.
package minio

import (
	"fmt"
	"reflect"

	"github.com/coreos/pkg/capnslog"
	opkit "github.com/rook/operator-kit"
	miniov1alpha1 "github.com/rook/rook/pkg/apis/minio.rook.io/v1alpha1"
	"github.com/rook/rook/pkg/clusterd"
	"github.com/rook/rook/pkg/operator/k8sutil"
	"k8s.io/api/apps/v1beta2"
	"k8s.io/api/core/v1"
	apiextensionsv1beta1 "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1beta1"
	"k8s.io/apimachinery/pkg/api/resource"
	meta_v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/tools/cache"
)

// TODO: A lot of these constants are specific to the KubeCon demo. Let's
// revisit these and determine what should be specified in the resource spec.
const (
	customResourceName       = "objectstore"
	customResourceNamePlural = "objectstores"
	minioCtrName             = "minio"
	minioLabel               = "minio"
	minioPVCName             = "minio-pvc"
	minioVolumeName          = "data"
	objectStoreDataDir       = "/data"
)

var logger = capnslog.NewPackageLogger("github.com/rook/rook", "minio-op-object")

// ObjectStoreResource represents the object store custom resource
var ObjectStoreResource = opkit.CustomResource{
	Name:    customResourceName,
	Plural:  customResourceNamePlural,
	Group:   miniov1alpha1.CustomResourceGroup,
	Version: miniov1alpha1.Version,
	Scope:   apiextensionsv1beta1.NamespaceScoped,
	Kind:    reflect.TypeOf(miniov1alpha1.ObjectStore{}).Name(),
}

// MinioController represents a controller object for object store custom resources
type MinioController struct {
	context   *clusterd.Context
	rookImage string
}

// NewMinioController create controller for watching object store custom resources created
func NewMinioController(context *clusterd.Context, rookImage string) *MinioController {
	return &MinioController{
		context:   context,
		rookImage: rookImage,
	}
}

// StartWatch watches for instances of ObjectStore custom resources and acts on them
func (c *MinioController) StartWatch(namespace string, stopCh chan struct{}) error {
	resourceHandlerFuncs := cache.ResourceEventHandlerFuncs{
		AddFunc:    c.onAdd,
		UpdateFunc: c.onUpdate,
		DeleteFunc: c.onDelete,
	}

	logger.Infof("start watching object store resources in namespace %s", namespace)
	watcher := opkit.NewWatcher(ObjectStoreResource, namespace, resourceHandlerFuncs, c.context.RookClientset.MinioV1alpha1().RESTClient())
	go watcher.Watch(&miniov1alpha1.ObjectStore{}, stopCh)

	return nil
}

func (c *MinioController) makeMinioHeadlessService(name, namespace string, spec miniov1alpha1.ObjectStoreSpec, ownerRef meta_v1.OwnerReference) (*v1.Service, error) {
	coreV1Client := c.context.Clientset.CoreV1()

	svc, err := coreV1Client.Services(namespace).Create(&v1.Service{
		ObjectMeta: meta_v1.ObjectMeta{
			Name:            name,
			Namespace:       namespace,
			Labels:          map[string]string{k8sutil.AppAttr: minioLabel},
			OwnerReferences: []meta_v1.OwnerReference{ownerRef},
		},
		Spec: v1.ServiceSpec{
			Selector:  map[string]string{k8sutil.AppAttr: minioLabel},
			Ports:     []v1.ServicePort{{Port: spec.Port}},
			ClusterIP: v1.ClusterIPNone,
		},
	})

	return svc, err
}

func (c *MinioController) buildMinioCtrArgs(statefulSetPrefix, headlessServiceName, namespace string, serverCount int32) []string {
	args := []string{"server"}
	for i := int32(0); i < serverCount; i++ {
		serverAddress := fmt.Sprintf("http://%s-%d.%s.%s%s", statefulSetPrefix, i, headlessServiceName, namespace, objectStoreDataDir)
		args = append(args, serverAddress)
	}

	logger.Infof("Building Minio container args: %v", args)
	return args
}

func (c *MinioController) makeMinioPodSpec(name, namespace string, ctrName string, ctrImage string, port int32, envVars map[string]string, numServers int32) v1.PodTemplateSpec {
	var env []v1.EnvVar
	for k, v := range envVars {
		env = append(env, v1.EnvVar{Name: k, Value: v})
	}

	podSpec := v1.PodTemplateSpec{
		ObjectMeta: meta_v1.ObjectMeta{
			Name:      name,
			Namespace: namespace,
			Labels:    map[string]string{k8sutil.AppAttr: minioLabel},
		},
		Spec: v1.PodSpec{
			Containers: []v1.Container{
				{
					Name:    ctrName,
					Image:   ctrImage,
					Env:     env,
					Command: []string{"/usr/bin/minio"},
					Ports:   []v1.ContainerPort{{ContainerPort: port}},
					Args:    c.buildMinioCtrArgs(name, name, namespace, numServers),
					VolumeMounts: []v1.VolumeMount{
						{
							Name:      minioVolumeName,
							MountPath: objectStoreDataDir,
						},
					},
				},
			},
			Volumes: []v1.Volume{
				{
					Name: minioVolumeName,
					VolumeSource: v1.VolumeSource{
						PersistentVolumeClaim: &v1.PersistentVolumeClaimVolumeSource{
							ClaimName: minioPVCName,
						},
					},
				},
			},
		},
	}

	return podSpec
}

func (c *MinioController) getAccessCredentials(secretName, namespace string) (string, string, error) {
	coreV1Client := c.context.Clientset.CoreV1()
	var getOpts meta_v1.GetOptions
	val, err := coreV1Client.Secrets(namespace).Get(secretName, getOpts)
	if err != nil {
		logger.Errorf("Unable to get secret with name=%s in namespace=%s: %v", secretName, namespace, err)
		return "", "", err
	}

	return string(val.Data["username"]), string(val.Data["password"]), nil
}

func validateObjectStoreSpec(spec miniov1alpha1.ObjectStoreSpec) error {
	// Verify node count.
	count := spec.Storage.NodeCount
	if count < 4 || count%2 != 0 {
		return fmt.Errorf("Node count must be greater than 3 and even.")
	}

	// Verify sane port.
	if spec.Port < 1024 {
		return fmt.Errorf("Invalid port %d", spec.Port)
	}

	return nil
}

func (c *MinioController) makeMinioStatefulSet(name, namespace string, spec miniov1alpha1.ObjectStoreSpec, ownerRef meta_v1.OwnerReference) (*v1beta2.StatefulSet, error) {
	appsClient := c.context.Clientset.AppsV1beta2()

	accessKey, secretKey, err := c.getAccessCredentials(spec.Credentials.Name, spec.Credentials.Namespace)
	if err != nil {
		return nil, err
	}

	envVars := map[string]string{
		"MINIO_ACCESS_KEY": accessKey,
		"MINIO_SECRET_KEY": secretKey,
	}

	podSpec := c.makeMinioPodSpec(name, namespace, minioCtrName, c.rookImage, spec.Port, envVars, int32(spec.Storage.NodeCount))

	nodeCount := int32(spec.Storage.NodeCount)
	ss := v1beta2.StatefulSet{
		ObjectMeta: meta_v1.ObjectMeta{
			Name:            name,
			Namespace:       namespace,
			Labels:          map[string]string{k8sutil.AppAttr: minioLabel},
			OwnerReferences: []meta_v1.OwnerReference{ownerRef},
		},
		Spec: v1beta2.StatefulSetSpec{
			Replicas: &nodeCount,
			Selector: &meta_v1.LabelSelector{
				MatchLabels: map[string]string{k8sutil.AppAttr: minioLabel},
			},
			Template: podSpec,
			VolumeClaimTemplates: []v1.PersistentVolumeClaim{
				{
					ObjectMeta: meta_v1.ObjectMeta{
						Name:      minioVolumeName,
						Namespace: namespace,
					},
					Spec: v1.PersistentVolumeClaimSpec{
						AccessModes: []v1.PersistentVolumeAccessMode{v1.ReadWriteOnce},
						Resources: v1.ResourceRequirements{
							Requests: v1.ResourceList{
								v1.ResourceStorage: resource.MustParse(spec.StorageSize),
							},
						},
					},
				},
			},
			ServiceName: name,
			// TODO: liveness probe
		},
	}

	return appsClient.StatefulSets(namespace).Create(&ss)
}

func (c *MinioController) onAdd(obj interface{}) {
	objectstore := obj.(*miniov1alpha1.ObjectStore).DeepCopy()

	ownerRef := meta_v1.OwnerReference{
		APIVersion: ObjectStoreResource.Version,
		Kind:       ObjectStoreResource.Kind,
		Name:       objectstore.Namespace,
		UID:        types.UID(objectstore.ObjectMeta.UID),
	}

	// Validate object store config.
	err := validateObjectStoreSpec(objectstore.Spec)
	if err != nil {
		logger.Errorf("failed to validate object store config")
		return
	}

	// Create the headless service.
	logger.Infof("Creating Minio headless service %s in namespace %s.", objectstore.Name, objectstore.Namespace)
	_, err = c.makeMinioHeadlessService(objectstore.Name, objectstore.Namespace, objectstore.Spec, ownerRef)
	if err != nil {
		logger.Errorf("failed to create minio headless service: %v", err)
		return
	}
	logger.Infof("Finished creating Minio headless service %s in namespace %s.", objectstore.Name, objectstore.Namespace)

	// Create the stateful set.
	logger.Infof("Creating Minio stateful set %s.", objectstore.Name)
	_, err = c.makeMinioStatefulSet(objectstore.Name, objectstore.Namespace, objectstore.Spec, ownerRef)
	if err != nil {
		logger.Errorf("failed to create minio stateful set: %v", err)
		return
	}
	logger.Infof("Finished creating Minio stateful set %s in namespace %s.", objectstore.Name, objectstore.Namespace)
}

func (c *MinioController) onUpdate(oldObj, newObj interface{}) {
	oldStore := oldObj.(*miniov1alpha1.ObjectStore).DeepCopy()
	newStore := newObj.(*miniov1alpha1.ObjectStore).DeepCopy()

	_ = oldStore
	_ = newStore

	logger.Infof("Received update on object store %s in namespace %s. This is currently unsupported.", oldStore.Name, oldStore.Namespace)
}

func (c *MinioController) onDelete(obj interface{}) {
	objectstore := obj.(*miniov1alpha1.ObjectStore).DeepCopy()
	logger.Infof("Delete Minio object store %s", objectstore.Name)

	// Cleanup is handled by the owner references set in 'onAdd' and the k8s garbage collector.
}
<|endoftext|>