go.mercari.io/datastore@v1.8.2/dsmiddleware/splitop/splitcall.go (about)

     1  package splitop
     2  
     3  import (
     4  	"context"
     5  
     6  	"go.mercari.io/datastore"
     7  )
     8  
     9  var _ datastore.Middleware = &splitHandler{}
    10  
    11  // New split call middleware will be returns.
    12  func New(opts ...Option) datastore.Middleware {
    13  	sh := &splitHandler{
    14  		putSplitThreshold: 500,
    15  		getSplitThreshold: 1000,
    16  	}
    17  	for _, opt := range opts {
    18  		opt.Apply(sh)
    19  	}
    20  	if sh.logf == nil {
    21  		sh.logf = func(ctx context.Context, format string, args ...interface{}) {}
    22  	}
    23  
    24  	return sh
    25  }
    26  
    27  // A Option is an option for splitop.
    28  type Option interface {
    29  	Apply(*splitHandler)
    30  }
    31  
    32  type splitHandler struct {
    33  	putSplitThreshold int
    34  	getSplitThreshold int
    35  
    36  	logf func(ctx context.Context, format string, args ...interface{})
    37  }
    38  
    39  func (sh *splitHandler) AllocateIDs(info *datastore.MiddlewareInfo, keys []datastore.Key) ([]datastore.Key, error) {
    40  	return info.Next.AllocateIDs(info, keys)
    41  }
    42  
    43  func (sh *splitHandler) PutMultiWithoutTx(info *datastore.MiddlewareInfo, keys []datastore.Key, psList []datastore.PropertyList) ([]datastore.Key, error) {
    44  	sh.logf(info.Context, "put %d keys", len(keys))
    45  	if sh.putSplitThreshold <= 0 || len(keys) <= sh.putSplitThreshold {
    46  		return info.Next.PutMultiWithoutTx(info, keys, psList)
    47  	}
    48  
    49  	retKeys := make([]datastore.Key, len(keys))
    50  	var mErr datastore.MultiError = make([]error, len(keys))
    51  	var foundErr bool
    52  	next := info.Next
    53  	for i := 0; i < len(keys); i += sh.putSplitThreshold {
    54  		end := i + sh.putSplitThreshold
    55  		if len(keys) < end {
    56  			end = len(keys)
    57  		}
    58  		sh.logf(info.Context, "put [%d, %d) range keys", i, end)
    59  		partialRetKeys, err := next.PutMultiWithoutTx(info, keys[i:end], psList[i:end])
    60  		for idx, key := range partialRetKeys {
    61  			retKeys[i+idx] = key
    62  		}
    63  		if mErr2, ok := err.(datastore.MultiError); ok {
    64  			for idx, err := range mErr2 {
    65  				if err != nil {
    66  					foundErr = true
    67  					mErr[i+idx] = err
    68  				}
    69  			}
    70  		} else if err != nil {
    71  			return nil, err
    72  		}
    73  	}
    74  	if foundErr {
    75  		return retKeys, mErr
    76  	}
    77  
    78  	return retKeys, nil
    79  }
    80  
    81  func (sh *splitHandler) PutMultiWithTx(info *datastore.MiddlewareInfo, keys []datastore.Key, psList []datastore.PropertyList) ([]datastore.PendingKey, error) {
    82  	return info.Next.PutMultiWithTx(info, keys, psList)
    83  }
    84  
    85  func (sh *splitHandler) GetMultiWithoutTx(info *datastore.MiddlewareInfo, keys []datastore.Key, psList []datastore.PropertyList) error {
    86  	sh.logf(info.Context, "get %d keys", len(keys))
    87  	if sh.getSplitThreshold <= 0 || len(keys) <= sh.getSplitThreshold {
    88  		return info.Next.GetMultiWithoutTx(info, keys, psList)
    89  	}
    90  	for len(psList) < len(keys) {
    91  		psList = append(psList, nil)
    92  	}
    93  
    94  	var mErr datastore.MultiError = make([]error, len(keys))
    95  	var foundErr bool
    96  	next := info.Next
    97  	for i := 0; i < len(keys); i += sh.getSplitThreshold {
    98  		end := i + sh.getSplitThreshold
    99  		if len(keys) < end {
   100  			end = len(keys)
   101  		}
   102  		sh.logf(info.Context, "get [%d, %d) range keys", i, end)
   103  		err := next.GetMultiWithoutTx(info, keys[i:end], psList[i:end])
   104  		if mErr2, ok := err.(datastore.MultiError); ok {
   105  			for idx, err := range mErr2 {
   106  				if err != nil {
   107  					foundErr = true
   108  					mErr[i+idx] = err
   109  				}
   110  			}
   111  		} else if err != nil {
   112  			return err
   113  		}
   114  	}
   115  	if foundErr {
   116  		return mErr
   117  	}
   118  
   119  	return nil
   120  }
   121  
   122  func (sh *splitHandler) GetMultiWithTx(info *datastore.MiddlewareInfo, keys []datastore.Key, psList []datastore.PropertyList) error {
   123  	return info.Next.GetMultiWithTx(info, keys, psList)
   124  }
   125  
   126  func (sh *splitHandler) DeleteMultiWithoutTx(info *datastore.MiddlewareInfo, keys []datastore.Key) error {
   127  	return info.Next.DeleteMultiWithoutTx(info, keys)
   128  }
   129  
   130  func (sh *splitHandler) DeleteMultiWithTx(info *datastore.MiddlewareInfo, keys []datastore.Key) error {
   131  	return info.Next.DeleteMultiWithTx(info, keys)
   132  }
   133  
   134  func (sh *splitHandler) PostCommit(info *datastore.MiddlewareInfo, tx datastore.Transaction, commit datastore.Commit) error {
   135  	return info.Next.PostCommit(info, tx, commit)
   136  }
   137  
   138  func (sh *splitHandler) PostRollback(info *datastore.MiddlewareInfo, tx datastore.Transaction) error {
   139  	return info.Next.PostRollback(info, tx)
   140  }
   141  
   142  func (sh *splitHandler) Run(info *datastore.MiddlewareInfo, q datastore.Query, qDump *datastore.QueryDump) datastore.Iterator {
   143  	return info.Next.Run(info, q, qDump)
   144  }
   145  
   146  func (sh *splitHandler) GetAll(info *datastore.MiddlewareInfo, q datastore.Query, qDump *datastore.QueryDump, psList *[]datastore.PropertyList) ([]datastore.Key, error) {
   147  	return info.Next.GetAll(info, q, qDump, psList)
   148  }
   149  
   150  func (sh *splitHandler) Next(info *datastore.MiddlewareInfo, q datastore.Query, qDump *datastore.QueryDump, iter datastore.Iterator, ps *datastore.PropertyList) (datastore.Key, error) {
   151  	return info.Next.Next(info, q, qDump, iter, ps)
   152  }
   153  
   154  func (sh *splitHandler) Count(info *datastore.MiddlewareInfo, q datastore.Query, qDump *datastore.QueryDump) (int, error) {
   155  	return info.Next.Count(info, q, qDump)
   156  }