github.com/tsuna/gohbase@v0.0.0-20250731002811-4ffcadfba63e/zk/client.go (about)

     1  // Copyright (C) 2015  The GoHBase Authors.  All rights reserved.
     2  // This file is part of GoHBase.
     3  // Use of this source code is governed by the Apache License 2.0
     4  // that can be found in the COPYING file.
     5  
     6  // Package zk encapsulates our interactions with ZooKeeper.
     7  package zk
     8  
     9  import (
    10  	"context"
    11  	"encoding/binary"
    12  	"fmt"
    13  	"log/slog"
    14  	"net"
    15  	"path"
    16  	"strings"
    17  	"time"
    18  
    19  	"github.com/go-zookeeper/zk"
    20  	"github.com/tsuna/gohbase/pb"
    21  	"google.golang.org/protobuf/proto"
    22  )
    23  
    24  type logger struct {
    25  	slogger *slog.Logger
    26  }
    27  
    28  func (l *logger) Printf(format string, args ...interface{}) {
    29  	l.slogger.Debug(fmt.Sprintf(format, args...))
    30  }
    31  
    32  func init() {
    33  	zk.DefaultLogger = &logger{slogger: slog.Default()}
    34  }
    35  
    36  // ResourceName is a type alias that is used to represent different resources
    37  // in ZooKeeper
    38  type ResourceName string
    39  
    40  // Prepend creates a new ResourceName with prefix prepended to the former ResourceName.
    41  func (r ResourceName) Prepend(prefix string) ResourceName {
    42  	return ResourceName(path.Join(prefix, string(r)))
    43  }
    44  
    45  const (
    46  	// Meta is a ResourceName that indicates that the location of the Meta
    47  	// table is what will be fetched
    48  	Meta = ResourceName("/meta-region-server")
    49  
    50  	// Master is a ResourceName that indicates that the location of the Master
    51  	// server is what will be fetched
    52  	Master = ResourceName("/master")
    53  )
    54  
    55  // Client is an interface of client that retrieves meta infomation from zookeeper
    56  type Client interface {
    57  	LocateResource(ResourceName) (string, error)
    58  }
    59  
    60  type client struct {
    61  	zks            []string
    62  	sessionTimeout time.Duration
    63  	dialer         func(ctx context.Context, network, addr string) (net.Conn, error)
    64  	logger         *slog.Logger
    65  }
    66  
    67  // NewClient establishes connection to zookeeper and returns the client
    68  func NewClient(zkquorum string, st time.Duration,
    69  	dialer func(ctx context.Context, network, addr string) (net.Conn, error),
    70  	slogger *slog.Logger) Client {
    71  
    72  	return &client{
    73  		zks:            strings.Split(zkquorum, ","),
    74  		sessionTimeout: st,
    75  		dialer:         dialer,
    76  		logger:         slogger,
    77  	}
    78  }
    79  
    80  // LocateResource returns address of the server for the specified resource.
    81  func (c *client) LocateResource(resource ResourceName) (string, error) {
    82  	var conn *zk.Conn
    83  	var err error
    84  	if c.dialer != nil {
    85  		conn, _, err = zk.Connect(c.zks, c.sessionTimeout, zk.WithDialer(makeZKDialer(c.dialer)))
    86  	} else {
    87  		conn, _, err = zk.Connect(c.zks, c.sessionTimeout)
    88  	}
    89  	if err != nil {
    90  		return "", fmt.Errorf("error connecting to ZooKeeper at %v: %s", c.zks, err)
    91  	}
    92  	defer conn.Close()
    93  
    94  	buf, _, err := conn.Get(string(resource))
    95  	if err != nil {
    96  		return "", fmt.Errorf("failed to read the %s znode: %s", resource, err)
    97  	}
    98  	if len(buf) == 0 {
    99  		panic(fmt.Errorf("%s was empty", resource))
   100  	} else if buf[0] != 0xFF {
   101  		return "", fmt.Errorf("the first byte of %s was 0x%x, not 0xFF", resource, buf[0])
   102  	}
   103  	metadataLen := binary.BigEndian.Uint32(buf[1:])
   104  	if metadataLen < 1 || metadataLen > 65000 {
   105  		return "", fmt.Errorf("invalid metadata length for %s: %d", resource, metadataLen)
   106  	}
   107  	buf = buf[1+4+metadataLen:]
   108  	magic := binary.BigEndian.Uint32(buf)
   109  	const pbufMagic = 1346524486 // 4 bytes: "PBUF"
   110  	if magic != pbufMagic {
   111  		return "", fmt.Errorf("invalid magic number for %s: %d", resource, magic)
   112  	}
   113  	buf = buf[4:]
   114  	var server *pb.ServerName
   115  	if resource == Meta {
   116  		meta := &pb.MetaRegionServer{}
   117  		err = proto.Unmarshal(buf, meta)
   118  		if err != nil {
   119  			return "",
   120  				fmt.Errorf("failed to deserialize the MetaRegionServer entry from ZK: %s", err)
   121  		}
   122  		server = meta.Server
   123  	} else {
   124  		master := &pb.Master{}
   125  		err = proto.Unmarshal(buf, master)
   126  		if err != nil {
   127  			return "",
   128  				fmt.Errorf("failed to deserialize the Master entry from ZK: %s", err)
   129  		}
   130  		server = master.Master
   131  	}
   132  	return net.JoinHostPort(*server.HostName, fmt.Sprint(*server.Port)), nil
   133  }
   134  
   135  func makeZKDialer(ctxDialer func(
   136  	ctx context.Context, network, addr string) (net.Conn, error)) zk.Dialer {
   137  	return func(network, addr string, timeout time.Duration) (net.Conn, error) {
   138  		ctx, cancel := context.WithTimeout(context.Background(), timeout)
   139  		defer cancel()
   140  		return ctxDialer(ctx, network, addr)
   141  	}
   142  }