github.com/erda-project/erda-infra@v1.0.9/providers/cassandra/writer.go (about)

     1  // Copyright (c) 2021 Terminus, Inc.
     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 cassandra
    16  
    17  import (
    18  	"time"
    19  
    20  	"github.com/erda-project/erda-infra/base/logs"
    21  	"github.com/gocql/gocql"
    22  )
    23  
    24  // StatementBuilder .
    25  type StatementBuilder interface {
    26  	GetStatement(data interface{}) (string, []interface{}, error)
    27  }
    28  
    29  type batchWriter struct {
    30  	session        *Session
    31  	builder        StatementBuilder
    32  	retry          int
    33  	retryDuration  time.Duration
    34  	log            logs.Logger
    35  	batchSizeBytes int
    36  }
    37  
    38  func (w *batchWriter) Write(data interface{}) error {
    39  	_, err := w.WriteN(data)
    40  	return err
    41  }
    42  
    43  func (w *batchWriter) WriteN(data ...interface{}) (int, error) {
    44  	if len(data) <= 0 {
    45  		return 0, nil
    46  	}
    47  	batchs := make([]*gocql.Batch, 0, 1)
    48  	sizeBytes := 0
    49  
    50  	batch := w.session.Session().NewBatch(gocql.LoggedBatch)
    51  	for _, item := range data {
    52  		stmt, args, err := w.builder.GetStatement(item)
    53  		if err != nil {
    54  			w.log.Errorf("fail to convert data to statement: %s", err)
    55  			continue
    56  		}
    57  		if w.batchSizeBytes > 0 {
    58  			sizeBytes += cqlSizeBytes(stmt, args)
    59  			if sizeBytes >= w.batchSizeBytes {
    60  				batchs = append(batchs, batch)
    61  				// reset
    62  				sizeBytes = 0
    63  				batch = w.session.Session().NewBatch(gocql.LoggedBatch)
    64  			}
    65  
    66  		}
    67  		batch.Query(stmt, args...)
    68  	}
    69  
    70  	if batch.Size() > 0 {
    71  		batchs = append(batchs, batch)
    72  	}
    73  
    74  	for _, batch := range batchs {
    75  		for i := 0; ; i++ {
    76  			err := w.session.Session().ExecuteBatch(batch)
    77  			if err != nil {
    78  				if w.retry == -1 || i < w.retry {
    79  					w.log.Warnf("fail to write batch(%d) to cassandra and retry after %s: %s", batch.Size(), w.retryDuration.String(), err)
    80  					time.Sleep(w.retryDuration)
    81  					continue
    82  				}
    83  				w.log.Errorf("fail to write batch(%d) to cassandra: %s", batch.Size(), err)
    84  				break
    85  			}
    86  			break
    87  		}
    88  	}
    89  
    90  	return batch.Size(), nil
    91  }
    92  
    93  func cqlSizeBytes(stmt string, args []interface{}) int {
    94  	size := len(stmt)
    95  	for _, item := range args {
    96  		switch v := item.(type) {
    97  		case string:
    98  			size += len(v)
    99  		case []byte:
   100  			size += len(v)
   101  		}
   102  	}
   103  	return size
   104  }
   105  
   106  func (w *batchWriter) Close() error {
   107  	if w.session == nil {
   108  		return nil
   109  	}
   110  	w.session.Close()
   111  	w.session = nil
   112  	return nil
   113  }