github.com/1aal/kubeblocks@v0.0.0-20231107070852-e1c03e598921/pkg/lorry/engines/etcd/manager.go (about)

     1  /*
     2  Copyright (C) 2022-2023 ApeCloud Co., Ltd
     3  
     4  This file is part of KubeBlocks project
     5  
     6  This program is free software: you can redistribute it and/or modify
     7  it under the terms of the GNU Affero General Public License as published by
     8  the Free Software Foundation, either version 3 of the License, or
     9  (at your option) any later version.
    10  
    11  This program is distributed in the hope that it will be useful
    12  but WITHOUT ANY WARRANTY; without even the implied warranty of
    13  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    14  GNU Affero General Public License for more details.
    15  
    16  You should have received a copy of the GNU Affero General Public License
    17  along with this program.  If not, see <http://www.gnu.org/licenses/>.
    18  */
    19  
    20  package etcd
    21  
    22  import (
    23  	"context"
    24  	"strconv"
    25  	"strings"
    26  	"time"
    27  
    28  	v3 "go.etcd.io/etcd/client/v3"
    29  	ctrl "sigs.k8s.io/controller-runtime"
    30  
    31  	"github.com/1aal/kubeblocks/pkg/lorry/engines"
    32  )
    33  
    34  const (
    35  	endpoint = "endpoint"
    36  
    37  	defaultPort        = 2379
    38  	defaultDialTimeout = 600 * time.Millisecond
    39  )
    40  
    41  type Manager struct {
    42  	engines.DBManagerBase
    43  	etcd     *v3.Client
    44  	endpoint string
    45  }
    46  
    47  var _ engines.DBManager = &Manager{}
    48  
    49  func NewManager(properties engines.Properties) (engines.DBManager, error) {
    50  	logger := ctrl.Log.WithName("ETCD")
    51  
    52  	managerBase, err := engines.NewDBManagerBase(logger)
    53  	if err != nil {
    54  		return nil, err
    55  	}
    56  
    57  	mgr := &Manager{
    58  		DBManagerBase: *managerBase,
    59  	}
    60  
    61  	var endpoints []string
    62  	endpoint, ok := properties[endpoint]
    63  	if ok {
    64  		mgr.endpoint = endpoint
    65  		endpoints = []string{endpoint}
    66  	}
    67  
    68  	cli, err := v3.New(v3.Config{
    69  		Endpoints:   endpoints,
    70  		DialTimeout: defaultDialTimeout,
    71  	})
    72  	if err != nil {
    73  		return nil, err
    74  	}
    75  
    76  	mgr.etcd = cli
    77  	return mgr, nil
    78  }
    79  
    80  func (mgr *Manager) IsDBStartupReady() bool {
    81  	if mgr.DBStartupReady {
    82  		return true
    83  	}
    84  
    85  	ctx, cancel := context.WithTimeout(context.Background(), defaultDialTimeout)
    86  	status, err := mgr.etcd.Status(ctx, mgr.endpoint)
    87  	cancel()
    88  	if err != nil {
    89  		mgr.Logger.Info("get etcd status failed", "error", err, "status", status)
    90  		return false
    91  	}
    92  
    93  	mgr.DBStartupReady = true
    94  	mgr.Logger.Info("DB startup ready")
    95  	return true
    96  }
    97  
    98  func (mgr *Manager) GetRunningPort() int {
    99  	index := strings.Index(mgr.endpoint, ":")
   100  	if index < 0 {
   101  		return defaultPort
   102  	}
   103  	port, err := strconv.Atoi(mgr.endpoint[index+1:])
   104  	if err != nil {
   105  		return defaultPort
   106  	}
   107  
   108  	return port
   109  }