vitess.io/vitess@v0.16.2/go/vt/vttablet/tabletserver/tx_prep_pool.go (about)

     1  /*
     2  Copyright 2019 The Vitess Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package tabletserver
    18  
    19  import (
    20  	"errors"
    21  	"fmt"
    22  	"sync"
    23  )
    24  
    25  var (
    26  	errPrepCommitting = errors.New("committing")
    27  	errPrepFailed     = errors.New("failed")
    28  )
    29  
    30  // TxPreparedPool manages connections for prepared transactions.
    31  // The Prepare functionality and associated orchestration
    32  // is done by TxPool.
    33  type TxPreparedPool struct {
    34  	mu       sync.Mutex
    35  	conns    map[string]*StatefulConnection
    36  	reserved map[string]error
    37  	capacity int
    38  }
    39  
    40  // NewTxPreparedPool creates a new TxPreparedPool.
    41  func NewTxPreparedPool(capacity int) *TxPreparedPool {
    42  	if capacity < 0 {
    43  		// If capacity is 0 all prepares will fail.
    44  		capacity = 0
    45  	}
    46  	return &TxPreparedPool{
    47  		conns:    make(map[string]*StatefulConnection, capacity),
    48  		reserved: make(map[string]error),
    49  		capacity: capacity,
    50  	}
    51  }
    52  
    53  // Put adds the connection to the pool. It returns an error
    54  // if the pool is full or on duplicate key.
    55  func (pp *TxPreparedPool) Put(c *StatefulConnection, dtid string) error {
    56  	pp.mu.Lock()
    57  	defer pp.mu.Unlock()
    58  	if _, ok := pp.reserved[dtid]; ok {
    59  		return errors.New("duplicate DTID in Prepare: " + dtid)
    60  	}
    61  	if _, ok := pp.conns[dtid]; ok {
    62  		return errors.New("duplicate DTID in Prepare: " + dtid)
    63  	}
    64  	if len(pp.conns) >= pp.capacity {
    65  		return fmt.Errorf("prepared transactions exceeded limit: %d", pp.capacity)
    66  	}
    67  	pp.conns[dtid] = c
    68  	return nil
    69  }
    70  
    71  // FetchForRollback returns the connection and removes it from the pool.
    72  // If the connection is not found, it returns nil. If the dtid
    73  // is in the reserved list, it means that an operator is trying
    74  // to resolve a previously failed commit. So, it removes the entry
    75  // and returns nil.
    76  func (pp *TxPreparedPool) FetchForRollback(dtid string) *StatefulConnection {
    77  	pp.mu.Lock()
    78  	defer pp.mu.Unlock()
    79  	if _, ok := pp.reserved[dtid]; ok {
    80  		delete(pp.reserved, dtid)
    81  		return nil
    82  	}
    83  	c := pp.conns[dtid]
    84  	delete(pp.conns, dtid)
    85  	return c
    86  }
    87  
    88  // FetchForCommit returns the connection for commit. Before returning,
    89  // it remembers the dtid in its reserved list as "committing". If
    90  // the dtid is already in the reserved list, it returns an error.
    91  // If the commit is successful, the dtid can be removed from the
    92  // reserved list by calling Forget. If the commit failed, SetFailed
    93  // must be called. This will inform future retries that the previous
    94  // commit failed.
    95  func (pp *TxPreparedPool) FetchForCommit(dtid string) (*StatefulConnection, error) {
    96  	pp.mu.Lock()
    97  	defer pp.mu.Unlock()
    98  	if err, ok := pp.reserved[dtid]; ok {
    99  		return nil, err
   100  	}
   101  	c, ok := pp.conns[dtid]
   102  	if ok {
   103  		delete(pp.conns, dtid)
   104  		pp.reserved[dtid] = errPrepCommitting
   105  	}
   106  	return c, nil
   107  }
   108  
   109  // SetFailed marks the reserved dtid as failed.
   110  // If there was no previous entry, one is created.
   111  func (pp *TxPreparedPool) SetFailed(dtid string) {
   112  	pp.mu.Lock()
   113  	defer pp.mu.Unlock()
   114  	pp.reserved[dtid] = errPrepFailed
   115  }
   116  
   117  // Forget removes the dtid from the reserved list.
   118  func (pp *TxPreparedPool) Forget(dtid string) {
   119  	pp.mu.Lock()
   120  	defer pp.mu.Unlock()
   121  	delete(pp.reserved, dtid)
   122  }
   123  
   124  // FetchAll removes all connections and returns them as a list.
   125  // It also forgets all reserved dtids.
   126  func (pp *TxPreparedPool) FetchAll() []*StatefulConnection {
   127  	pp.mu.Lock()
   128  	defer pp.mu.Unlock()
   129  	conns := make([]*StatefulConnection, 0, len(pp.conns))
   130  	for _, c := range pp.conns {
   131  		conns = append(conns, c)
   132  	}
   133  	pp.conns = make(map[string]*StatefulConnection, pp.capacity)
   134  	pp.reserved = make(map[string]error)
   135  	return conns
   136  }