github.com/matrixorigin/matrixone@v1.2.0/pkg/vm/engine/tae/logstore/driver/logservicedriver/clientpool.go (about)

     1  // Copyright 2021 Matrix Origin
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package logservicedriver
    16  
    17  import (
    18  	"sync"
    19  	"sync/atomic"
    20  	"time"
    21  
    22  	"github.com/matrixorigin/matrixone/pkg/common/moerr"
    23  	"github.com/matrixorigin/matrixone/pkg/logservice"
    24  )
    25  
    26  var ErrNoClientAvailable = moerr.NewInternalErrorNoCtx("no client available")
    27  var ErrClientPoolClosed = moerr.NewInternalErrorNoCtx("client pool closed")
    28  
    29  type clientConfig struct {
    30  	cancelDuration        time.Duration
    31  	recordSize            int
    32  	clientFactory         LogServiceClientFactory
    33  	GetClientRetryTimeOut time.Duration
    34  	retryDuration         time.Duration
    35  }
    36  
    37  type clientWithRecord struct {
    38  	c      logservice.Client
    39  	record logservice.LogRecord
    40  	id     int
    41  }
    42  
    43  func newClient(factory LogServiceClientFactory, recordSize int, retryDuration time.Duration) *clientWithRecord {
    44  	logserviceClient, err := factory()
    45  	if err != nil {
    46  		RetryWithTimeout(retryDuration, func() (shouldReturn bool) {
    47  			logserviceClient, err = factory()
    48  			return err == nil
    49  		})
    50  		if err != nil {
    51  			panic(err)
    52  		}
    53  	}
    54  	c := &clientWithRecord{
    55  		c:      logserviceClient,
    56  		record: logserviceClient.GetLogRecord(recordSize),
    57  	}
    58  	return c
    59  }
    60  
    61  func (c *clientWithRecord) Close() {
    62  	c.c.Close()
    63  }
    64  
    65  func (c *clientWithRecord) TryResize(size int) {
    66  	if len(c.record.Payload()) < size {
    67  		c.record = c.c.GetLogRecord(size)
    68  	}
    69  }
    70  
    71  type clientpool struct {
    72  	maxCount   int
    73  	count      int
    74  	getTimeout time.Duration
    75  
    76  	closed atomic.Int32
    77  
    78  	freeClients   []*clientWithRecord
    79  	clientFactory func() *clientWithRecord
    80  	closefn       func(*clientWithRecord)
    81  	mu            sync.Mutex
    82  }
    83  
    84  func newClientPool(maxsize int, cfg *clientConfig) *clientpool {
    85  	pool := &clientpool{
    86  		maxCount:    maxsize,
    87  		getTimeout:  cfg.GetClientRetryTimeOut,
    88  		freeClients: make([]*clientWithRecord, maxsize),
    89  		mu:          sync.Mutex{},
    90  	}
    91  	pool.clientFactory = pool.newClientFactory(cfg)
    92  	pool.closefn = pool.onClose
    93  
    94  	for i := 0; i < maxsize; i++ {
    95  		pool.freeClients[i] = pool.clientFactory()
    96  	}
    97  	return pool
    98  }
    99  
   100  func (c *clientpool) newClientFactory(cfg *clientConfig) func() *clientWithRecord {
   101  	return func() *clientWithRecord {
   102  		c.count++
   103  		client := newClient(cfg.clientFactory, cfg.recordSize, cfg.retryDuration)
   104  		client.id = c.count
   105  		return client
   106  	}
   107  }
   108  
   109  func (c *clientpool) Close() {
   110  	if c.IsClosed() {
   111  		return
   112  	}
   113  	c.mu.Lock()
   114  	if c.IsClosed() {
   115  		c.mu.Unlock()
   116  		return
   117  	}
   118  	for _, client := range c.freeClients {
   119  		c.closefn(client)
   120  	}
   121  	c.closed.Store(1)
   122  	c.mu.Unlock()
   123  }
   124  
   125  func (c *clientpool) onClose(client *clientWithRecord) {
   126  	client.Close()
   127  }
   128  
   129  // func (c *clientpool) GetAndWait() (*clientWithRecord, error) {
   130  // 	client, err := c.Get()
   131  // 	if err == ErrNoClientAvailable {
   132  // 		retryWithTimeout(c.getTimeout, func() (shouldReturn bool) {
   133  // 			client, err = c.Get()
   134  // 			return err != ErrNoClientAvailable
   135  // 		})
   136  // 	}
   137  // 	return client, nil
   138  // }
   139  
   140  func (c *clientpool) Get() (*clientWithRecord, error) {
   141  	if c.IsClosed() {
   142  		return nil, ErrClientPoolClosed
   143  	}
   144  	c.mu.Lock()
   145  	defer c.mu.Unlock()
   146  	if c.IsClosed() {
   147  		return nil, ErrClientPoolClosed
   148  	}
   149  	if len(c.freeClients) == 0 {
   150  		if c.count == c.maxCount {
   151  			return nil, ErrNoClientAvailable
   152  		}
   153  		return c.clientFactory(), nil
   154  	}
   155  	client := c.freeClients[len(c.freeClients)-1]
   156  	c.freeClients = c.freeClients[:len(c.freeClients)-1]
   157  	return client, nil
   158  }
   159  
   160  func (c *clientpool) IsClosed() bool {
   161  	return c.closed.Load() == 1
   162  }
   163  
   164  func (c *clientpool) Put(client *clientWithRecord) {
   165  	if c.IsClosed() {
   166  		c.closefn(client)
   167  		return
   168  	}
   169  	c.mu.Lock()
   170  	if c.IsClosed() {
   171  		c.closefn(client)
   172  		c.mu.Unlock()
   173  		return
   174  	}
   175  	c.count--
   176  	c.freeClients = append(c.freeClients, client)
   177  	defer c.mu.Unlock()
   178  }