github.phpd.cn/hashicorp/consul@v1.4.5/agent/consul/state/autopilot.go (about)

     1  package state
     2  
     3  import (
     4  	"fmt"
     5  
     6  	"github.com/hashicorp/consul/agent/consul/autopilot"
     7  	"github.com/hashicorp/go-memdb"
     8  )
     9  
    10  // autopilotConfigTableSchema returns a new table schema used for storing
    11  // the autopilot configuration
    12  func autopilotConfigTableSchema() *memdb.TableSchema {
    13  	return &memdb.TableSchema{
    14  		Name: "autopilot-config",
    15  		Indexes: map[string]*memdb.IndexSchema{
    16  			"id": &memdb.IndexSchema{
    17  				Name:         "id",
    18  				AllowMissing: true,
    19  				Unique:       true,
    20  				Indexer: &memdb.ConditionalIndex{
    21  					Conditional: func(obj interface{}) (bool, error) { return true, nil },
    22  				},
    23  			},
    24  		},
    25  	}
    26  }
    27  
    28  func init() {
    29  	registerSchema(autopilotConfigTableSchema)
    30  }
    31  
    32  // Autopilot is used to pull the autopilot config from the snapshot.
    33  func (s *Snapshot) Autopilot() (*autopilot.Config, error) {
    34  	c, err := s.tx.First("autopilot-config", "id")
    35  	if err != nil {
    36  		return nil, err
    37  	}
    38  
    39  	config, ok := c.(*autopilot.Config)
    40  	if !ok {
    41  		return nil, nil
    42  	}
    43  
    44  	return config, nil
    45  }
    46  
    47  // Autopilot is used when restoring from a snapshot.
    48  func (s *Restore) Autopilot(config *autopilot.Config) error {
    49  	if err := s.tx.Insert("autopilot-config", config); err != nil {
    50  		return fmt.Errorf("failed restoring autopilot config: %s", err)
    51  	}
    52  
    53  	return nil
    54  }
    55  
    56  // AutopilotConfig is used to get the current Autopilot configuration.
    57  func (s *Store) AutopilotConfig() (uint64, *autopilot.Config, error) {
    58  	tx := s.db.Txn(false)
    59  	defer tx.Abort()
    60  
    61  	// Get the autopilot config
    62  	c, err := tx.First("autopilot-config", "id")
    63  	if err != nil {
    64  		return 0, nil, fmt.Errorf("failed autopilot config lookup: %s", err)
    65  	}
    66  
    67  	config, ok := c.(*autopilot.Config)
    68  	if !ok {
    69  		return 0, nil, nil
    70  	}
    71  
    72  	return config.ModifyIndex, config, nil
    73  }
    74  
    75  // AutopilotSetConfig is used to set the current Autopilot configuration.
    76  func (s *Store) AutopilotSetConfig(idx uint64, config *autopilot.Config) error {
    77  	tx := s.db.Txn(true)
    78  	defer tx.Abort()
    79  
    80  	if err := s.autopilotSetConfigTxn(idx, tx, config); err != nil {
    81  		return err
    82  	}
    83  
    84  	tx.Commit()
    85  	return nil
    86  }
    87  
    88  // AutopilotCASConfig is used to try updating the Autopilot configuration with a
    89  // given Raft index. If the CAS index specified is not equal to the last observed index
    90  // for the config, then the call is a noop,
    91  func (s *Store) AutopilotCASConfig(idx, cidx uint64, config *autopilot.Config) (bool, error) {
    92  	tx := s.db.Txn(true)
    93  	defer tx.Abort()
    94  
    95  	// Check for an existing config
    96  	existing, err := tx.First("autopilot-config", "id")
    97  	if err != nil {
    98  		return false, fmt.Errorf("failed autopilot config lookup: %s", err)
    99  	}
   100  
   101  	// If the existing index does not match the provided CAS
   102  	// index arg, then we shouldn't update anything and can safely
   103  	// return early here.
   104  	e, ok := existing.(*autopilot.Config)
   105  	if !ok || e.ModifyIndex != cidx {
   106  		return false, nil
   107  	}
   108  
   109  	if err := s.autopilotSetConfigTxn(idx, tx, config); err != nil {
   110  		return false, err
   111  	}
   112  
   113  	tx.Commit()
   114  	return true, nil
   115  }
   116  
   117  func (s *Store) autopilotSetConfigTxn(idx uint64, tx *memdb.Txn, config *autopilot.Config) error {
   118  	// Check for an existing config
   119  	existing, err := tx.First("autopilot-config", "id")
   120  	if err != nil {
   121  		return fmt.Errorf("failed autopilot config lookup: %s", err)
   122  	}
   123  
   124  	// Set the indexes.
   125  	if existing != nil {
   126  		config.CreateIndex = existing.(*autopilot.Config).CreateIndex
   127  	} else {
   128  		config.CreateIndex = idx
   129  	}
   130  	config.ModifyIndex = idx
   131  
   132  	if err := tx.Insert("autopilot-config", config); err != nil {
   133  		return fmt.Errorf("failed updating autopilot config: %s", err)
   134  	}
   135  	return nil
   136  }