graph.go

Documentation: edge-infra.dev/pkg/edge/rollouts

     1  package rollouts
     2  
     3  import (
     4  	"encoding/json"
     5  	"fmt"
     6  	"reflect"
     7  	"strings"
     8  )
     9  
    10  type RolloutState string
    11  
    12  const (
    13  	RolloutPending    RolloutState = "pending"
    14  	RolloutInProgress RolloutState = "in_progress"
    15  	RolloutComplete   RolloutState = "complete"
    16  	RolloutError      RolloutState = "error"
    17  )
    18  
    19  // NodeKey is a key to look up a node within a graph. It must be unique within a graph
    20  type NodeKey string
    21  
    22  // NodeExecutionResult TODO(dk185217): can be extended to a struct type later if necessary
    23  // type NodeExecutionResult bool
    24  
    25  type NodeExecutionResult struct {
    26  	Key     NodeKey
    27  	State   NodeState
    28  	Message string
    29  	// TODO(dk185217): this signals the status of the overall rollout
    30  	// - possibly move somewhere else. here or at a higher level?
    31  	RolloutState RolloutState
    32  }
    33  
    34  // Done is ... TODO(dk185217): can be extended to a struct type later if necessary
    35  func (r NodeExecutionResult) Done() bool {
    36  	switch r.State {
    37  	case Complete:
    38  		return true
    39  	default:
    40  		return false
    41  	}
    42  }
    43  
    44  func (r NodeExecutionResult) String() string {
    45  	return r.Message
    46  }
    47  
    48  // RolloutGraphEdge represents an edge in the graph. A graph is defined by a list of its edges
    49  type RolloutGraphEdge struct {
    50  	From NodeKey
    51  	To   NodeKey
    52  }
    53  
    54  type NodeMap map[NodeKey]RolloutGraphNode
    55  
    56  func (nm NodeMap) AddNode(node RolloutGraphNode) {
    57  	nm[node.GetKey()] = node
    58  }
    59  
    60  // RolloutGraph is an instance of a rollout plan
    61  type RolloutGraph struct {
    62  	// TODO(dk185217): consider un-exporting these fields since we dont directly unmarshal into them
    63  	ID                   string
    64  	Nodes                NodeMap
    65  	Edges                []RolloutGraphEdge
    66  	Current              []NodeKey
    67  	NodeExecutionResults map[NodeKey]NodeExecutionResult
    68  
    69  	opts *options
    70  }
    71  
    72  // RolloutGraphJSON is the intermediate format used to implement custom json unmarshaling
    73  type RolloutGraphJSON struct {
    74  	ID                   string                          `json:"id,omitempty"`
    75  	Current              []NodeKey                       `json:"current"`
    76  	Nodes                []BaseNodeJSON                  `json:"nodes"`
    77  	Edges                []RolloutGraphEdge              `json:"edges"`
    78  	NodeExecutionResults map[NodeKey]NodeExecutionResult `json:"nodeExecutionResults"`
    79  	GraphState           string                          `json:"graph_state,omitempty"`
    80  }
    81  
    82  func NewRolloutGraphFromJSON(data json.RawMessage, opts ...Option) (*RolloutGraph, error) {
    83  	graph := &RolloutGraph{
    84  		opts: makeOptions(opts...),
    85  	}
    86  	err := json.Unmarshal(data, graph)
    87  	if err != nil {
    88  		return nil, err
    89  	}
    90  	return graph, nil
    91  }
    92  
    93  // RolloutGraph ...
    94  func NewRolloutGraph(plan RolloutPlan, opts ...Option) *RolloutGraph {
    95  	// TODO(https://github.com/ncr-swt-retail/edge-roadmap/issues/10936): Parameter 'plan' is unused.
    96  	// This is where we get the GraphConfiguration to create a graph for the first time. Ie, unmarshal
    97  	// the plan.GraphConfiguration into a RolloutGraph and initialize its state:
    98  	// - when first created, all nodes should be Pending, and the current node will be a StartNode
    99  	// - there will be no NodeExecutionResults yet, because nothing has been executed
   100  	// But, now it is a valid RolloutGraph that can be Marshaled/Unmarshaled at will
   101  	options := makeOptions(opts...)
   102  
   103  	g := &RolloutGraph{
   104  		Current:              plan.Initial,
   105  		NodeExecutionResults: map[NodeKey]NodeExecutionResult{},
   106  	}
   107  
   108  	for _, node := range plan.Nodes {
   109  		// Initialize node states since plans are stateless
   110  		node.initState()
   111  
   112  		// Set per node graph initialization values
   113  		switch n := node.(type) {
   114  		case *TimerGate:
   115  			if options.timeSource != nil {
   116  				n.time = options.timeSource
   117  			} else {
   118  				n.time = defaultTimeSource{}
   119  			}
   120  		}
   121  	}
   122  
   123  	g.Nodes = plan.Nodes
   124  	g.Edges = plan.Edges
   125  
   126  	return g
   127  }
   128  
   129  func (g *RolloutGraph) String() string {
   130  	if len(g.Current) == 0 {
   131  		return "{ Current: [] }"
   132  	}
   133  	currents := []string{}
   134  	for _, currKey := range g.Current {
   135  		curr := g.Nodes[currKey]
   136  		currents = append(currents, fmt.Sprintf("%s (%s)", curr.GetLabel(), curr.GetKey()))
   137  	}
   138  	currentsStr := strings.Join(currents, ",")
   139  	// TODO: Expand - Printf calls are going to look empty
   140  	return fmt.Sprintf("{ Current: [%v] }", currentsStr)
   141  }
   142  
   143  // UnmarshalJSON is ...
   144  func (g *RolloutGraph) UnmarshalJSON(data []byte) error {
   145  	// Unmarshal raw json payload into intermediate type
   146  	rolloutJSON := &RolloutGraphJSON{}
   147  	err := json.Unmarshal(data, rolloutJSON)
   148  	if err != nil {
   149  		return fmt.Errorf("failed to unmarshal rollout graph payload: %v", err)
   150  	}
   151  
   152  	// ID
   153  	g.ID = rolloutJSON.ID
   154  
   155  	// Nodes. Unmarshal nodes into concrete types
   156  	g.Nodes = make(map[NodeKey]RolloutGraphNode, len(rolloutJSON.Nodes))
   157  	for _, n := range rolloutJSON.Nodes {
   158  		node, err := g.UnmarshalRolloutGraphNode(n)
   159  		if err != nil {
   160  			return err
   161  		}
   162  		g.Nodes[node.GetKey()] = node
   163  	}
   164  
   165  	g.Edges = make([]RolloutGraphEdge, len(rolloutJSON.Edges))
   166  	for i, edge := range rolloutJSON.Edges {
   167  		// For each edge, create links between the respective nodes
   168  		g.Edges[i] = RolloutGraphEdge{To: edge.To, From: edge.From}
   169  		toNode, found := g.Nodes[edge.To]
   170  		if !found {
   171  			return fmt.Errorf("node key %s was stated in edge but not in nodes", edge.To)
   172  		}
   173  		fromNode, found := g.Nodes[edge.From]
   174  		if !found {
   175  			return fmt.Errorf("node key %s was stated in edge but not in nodes", edge.From)
   176  		}
   177  		toNode.AddDependency(fromNode)
   178  		fromNode.AddNext(toNode)
   179  	}
   180  
   181  	// Current. Rote copying
   182  	g.Current = rolloutJSON.Current
   183  
   184  	// ExecutionResults. Rote copying
   185  	g.NodeExecutionResults = make(map[NodeKey]NodeExecutionResult, len(rolloutJSON.NodeExecutionResults))
   186  	for key, result := range rolloutJSON.NodeExecutionResults {
   187  		g.NodeExecutionResults[key] = result
   188  	}
   189  
   190  	return nil
   191  }
   192  
   193  // UnmarshalRolloutGraphNode unmarshals JSON into nodes. Some defaults are set
   194  // per node type to accommodate if the node comes from a Plan's JSON
   195  func (g *RolloutGraph) UnmarshalRolloutGraphNode(n BaseNodeJSON) (RolloutGraphNode, error) {
   196  	switch n.NodeType {
   197  	case "TargetGroup":
   198  		tg := &TargetGroup{}
   199  		if err := json.Unmarshal(n.Data, tg); err != nil {
   200  			return nil, fmt.Errorf("failed to unmarshal node: %v", err)
   201  		}
   202  		return tg, nil
   203  	case "TimerGate":
   204  		tg := &TimerGate{}
   205  		if err := json.Unmarshal(n.Data, tg); err != nil {
   206  			return nil, fmt.Errorf("failed to unmarshal node: %v", err)
   207  		}
   208  		tg.time = g.opts.timeSource
   209  		return tg, nil
   210  	case "ApprovalGate":
   211  		ag := &ApprovalGate{}
   212  		if err := json.Unmarshal(n.Data, ag); err != nil {
   213  			return nil, fmt.Errorf("failed to unmarshal node: %v", err)
   214  		}
   215  		return ag, nil
   216  	case "ScheduleGate":
   217  		sg := &ScheduleGate{}
   218  		if err := json.Unmarshal(n.Data, sg); err != nil {
   219  			return nil, fmt.Errorf("failed to unmarshal node: %v", err)
   220  		}
   221  		sg.time = g.opts.timeSource
   222  		return sg, nil
   223  	default:
   224  		return nil, fmt.Errorf("failed to unmarshal node, unknown NodeType: %s", n.NodeType)
   225  	}
   226  }
   227  
   228  func (g *RolloutGraph) MarshalJSON() ([]byte, error) {
   229  	// Current. Copy all current node keys
   230  	current := make([]NodeKey, len(g.Current))
   231  	for i, currKey := range g.Current {
   232  		curr := g.Nodes[currKey]
   233  		current[i] = curr.GetKey()
   234  	}
   235  
   236  	// Nodes. Convert all nodes to serializable types
   237  	nodes := make([]BaseNodeJSON, len(g.Nodes))
   238  	i := 0
   239  	for _, node := range g.Nodes {
   240  		data, err := json.Marshal(node)
   241  		if err != nil {
   242  			return nil, fmt.Errorf("failed to marshal node: %v", err)
   243  		}
   244  		baseNodeJSON := BaseNodeJSON{
   245  			// Seems to be necessary for pointer types, eg *TargetGroup -> TargetGroup. If this is a hassle to maintain
   246  			// or causes issues, a GetType func could be added to the node interface and implemented by each node type
   247  			NodeType: reflect.Indirect(reflect.ValueOf(node)).Type().Name(),
   248  			Data:     data,
   249  		}
   250  		nodes[i] = baseNodeJSON
   251  		i++
   252  	}
   253  
   254  	// Edges. Rote copying
   255  	edges := make([]RolloutGraphEdge, len(g.Edges))
   256  	copy(edges, g.Edges)
   257  
   258  	// ExecutionResults. Rote copying
   259  	results := map[NodeKey]NodeExecutionResult{}
   260  	for key, result := range g.NodeExecutionResults {
   261  		results[key] = result
   262  	}
   263  
   264  	jsonData := &RolloutGraphJSON{
   265  		Current:              current,
   266  		Nodes:                nodes,
   267  		NodeExecutionResults: results,
   268  		Edges:                edges,
   269  	}
   270  
   271  	return json.Marshal(jsonData)
   272  }
   273  
   274  func ConnectEdge(from RolloutGraphNode, to RolloutGraphNode, edges []RolloutGraphEdge) []RolloutGraphEdge {
   275  	from.AddNext(to)
   276  	to.AddDependency(from)
   277  	edge := RolloutGraphEdge{From: from.GetKey(), To: to.GetKey()}
   278  	edges = append(edges, edge)
   279  	return edges
   280  }
   281
View as plain text