state.go

package httpsched

import (
	"fmt"
	"log"
	"net/http"
	"sync"

	"github.com/mesos/mesos-go/api/v1/lib"
	"github.com/mesos/mesos-go/api/v1/lib/encoding"
	"github.com/mesos/mesos-go/api/v1/lib/httpcli"
	"github.com/mesos/mesos-go/api/v1/lib/httpcli/apierrors"
	"github.com/mesos/mesos-go/api/v1/lib/scheduler"
	"github.com/mesos/mesos-go/api/v1/lib/scheduler/calls"
)

const (
	headerMesosStreamID = "Mesos-Stream-Id"
	debug               = false
)

type StateError string

func (err StateError) Error() string { return string(err) }

var (
	errMissingMesosStreamId = httpcli.ProtocolError("missing Mesos-Stream-Id header expected with successful SUBSCRIBE")
	errAlreadySubscribed    = StateError("already subscribed, cannot re-issue a SUBSCRIBE call")
)

type (
	// state implements calls.Caller and tracks connectivity with Mesos
	state struct {
		client *client // client is a handle to the original underlying HTTP client

		m      sync.Mutex
		fn     stateFn      // fn is the next state function to execute
		caller calls.Caller // caller is (maybe) used by a state function to execute a call

		call *scheduler.Call // call is the next call to execute
		resp mesos.Response  // resp is the Mesos response from the most recently executed call
		err  error           // err is the error from the most recently executed call
	}

	stateFn func(*state) stateFn
)

func maybeLogged(f httpcli.DoFunc) httpcli.DoFunc {
	if debug {
		return func(req *http.Request) (*http.Response, error) {
			log.Println("wrapping request", req.URL, req.Header)
			resp, err := f(req)
			if err == nil {
				log.Printf("status %d", resp.StatusCode)
				for k := range resp.Header {
					log.Println("header " + k + ": " + resp.Header.Get(k))
				}
			}
			return resp, err
		}
	}
	return f
}

// DisconnectionDetector is a programmable response decorator that attempts to detect errors
// that should transition the state from "connected" to "disconnected". Detector implementations
// are expected to invoke the `disconnect` callback in order to initiate the disconnection.
//
// The default implementation will transition to a disconnected state when:
//   - an error occurs while decoding an object from the subscription stream
//   - mesos reports an ERROR-type scheduler.Event object via the subscription stream
//   - an object on the stream does not decode to a *scheduler.Event (sanity check)
//
// Consumers of this package may choose to override default behavior by overwriting the default
// value of this var, but should exercise caution: failure to properly transition to a disconnected
// state may cause subsequent Call operations to fail (without recourse).
var DisconnectionDetector = func(disconnect func()) mesos.ResponseDecorator {
	return mesos.ResponseDecoratorFunc(func(resp mesos.Response) mesos.Response {
		return &mesos.ResponseWrapper{
			Response: resp,
			Decoder:  disconnectionDecoder(resp, disconnect),
		}
	})
}

func disconnectionDecoder(decoder encoding.Decoder, disconnect func()) encoding.Decoder {
	return encoding.DecoderFunc(func(u encoding.Unmarshaler) (err error) {
		err = decoder.Decode(u)
		if err != nil {
			disconnect()
			return
		}
		switch e := u.(type) {
		case (*scheduler.Event):
			if e.GetType() == scheduler.Event_ERROR {
				// the mesos scheduler API recommends that scheduler implementations
				// resubscribe in this case. we initiate the disconnection here because
				// it is assumed to be convenient for most framework implementations.
				disconnect()
			}
		default:
			// sanity check: this should never happen in practice.
			err = httpcli.ProtocolError(
				fmt.Sprintf("unexpected object on subscription event stream: %v", e))
			disconnect()
		}
		return
	})
}

func disconnectedFn(state *state) stateFn {
	// (a) validate call = SUBSCRIBE
	if state.call.GetType() != scheduler.Call_SUBSCRIBE {
		state.resp = nil
		state.err = &apierrors.Error{Code: apierrors.CodeUnsubscribed}
		return disconnectedFn
	}

	// (b) prepare client for a subscription call
	var (
		mesosStreamID = ""
		undoable      = httpcli.WrapDoer(func(f httpcli.DoFunc) httpcli.DoFunc {
			f = maybeLogged(f)
			return func(req *http.Request) (resp *http.Response, err error) {
				resp, err = f(req)
				if err == nil && resp.StatusCode == 200 {
					// grab Mesos-Stream-Id header; if missing then
					// close the response body and return an error
					mesosStreamID = resp.Header.Get(headerMesosStreamID)
					if mesosStreamID == "" {
						resp.Body.Close()
						resp = nil
						err = errMissingMesosStreamId
					}
				}
				return
			}
		})
		subscribeCaller = &callerTemporary{
			opt:            undoable,
			callerInternal: state.client,
			requestOpts:    []httpcli.RequestOpt{httpcli.Close(true)},
		}
	)

	// (c) execute the call, save the result in resp, err
	stateResp, stateErr := subscribeCaller.Call(state.call)
	state.err = stateErr

	// (d) if err != nil return disconnectedFn since we're unsubscribed
	if stateErr != nil {
		if stateResp != nil {
			stateResp.Close()
		}
		state.resp = nil
		return disconnectedFn
	}

	transitionToDisconnected := func() {
		state.m.Lock()
		defer state.m.Unlock()
		state.fn = disconnectedFn
		_ = stateResp.Close() // swallow any error here
	}

	// wrap the response: any errors processing the subscription stream should result in a
	// transition to a disconnected state ASAP.
	state.resp = DisconnectionDetector(transitionToDisconnected).Decorate(stateResp)

	// (e) else prepare callerTemporary w/ special header, return connectedFn since we're now subscribed
	state.caller = &callerTemporary{
		opt:            httpcli.DefaultHeader(headerMesosStreamID, mesosStreamID),
		callerInternal: state.client,
	}
	return connectedFn
}

// CodesIndicatingSubscriptionLoss is a set of apierror.Code entries which each indicate that
// the event subscription stream has been severed between the scheduler and mesos. It's respresented
// as a public map variable so that clients can program additional error codes (if such are discovered)
// without hacking the code of the mesos-go library directly.
var CodesIndicatingSubscriptionLoss = func(codes ...apierrors.Code) map[apierrors.Code]struct{} {
	result := make(map[apierrors.Code]struct{}, len(codes))
	for _, code := range codes {
		result[code] = struct{}{}
	}
	return result
}(
	// expand this list as we discover other errors that guarantee we've lost our event subscription.
	apierrors.CodeUnsubscribed,
)

func errorIndicatesSubscriptionLoss(err error) (result bool) {
	if apiError, ok := err.(*apierrors.Error); ok {
		_, result = CodesIndicatingSubscriptionLoss[apiError.Code]
	}
	// TODO(jdef) should other error types be considered here as well?
	return
}

func connectedFn(state *state) stateFn {
	// (a) validate call != SUBSCRIBE
	if state.call.GetType() == scheduler.Call_SUBSCRIBE {
		state.resp = nil

		// TODO(jdef) not super happy with this error: I don't think that mesos minds if we issue
		// redundant subscribe calls. However, the state tracking mechanism in this module can't
		// cope with it (e.g. we'll need to track a new stream-id, etc).
		// We make a best effort to transition to a disconnected state if we detect protocol errors,
		// error events, or mesos-generated "not subscribed" errors. But we don't handle things such
		// as, for example, authentication errors. Granted, the initial subscribe call should fail
		// if authentication is an issue, so we should never end up here. I'm not convinced there's
		// not other edge cases though with respect to other error codes.
		state.err = errAlreadySubscribed
		return connectedFn
	}

	// (b) execute call, save the result in resp, err
	state.resp, state.err = state.caller.Call(state.call)

	if errorIndicatesSubscriptionLoss(state.err) {
		// properly transition back to a disconnected state if mesos thinks that we're unsubscribed
		return disconnectedFn
	}

	// stay connected, don't attempt to interpret other errors here
	return connectedFn
}

func (state *state) Call(call *scheduler.Call) (resp mesos.Response, err error) {
	state.m.Lock()
	defer state.m.Unlock()
	state.call = call
	state.fn = state.fn(state)

	if debug && state.err != nil {
		log.Print(*call, state.err)
	}

	return state.resp, state.err
}