github.com/cilium/statedb@v0.3.2/derive.go (about)

     1  // SPDX-License-Identifier: Apache-2.0
     2  // Copyright Authors of Cilium
     3  
     4  package statedb
     5  
     6  import (
     7  	"context"
     8  
     9  	"github.com/cilium/hive/cell"
    10  	"github.com/cilium/hive/job"
    11  )
    12  
    13  type DeriveResult int
    14  
    15  const (
    16  	DeriveInsert DeriveResult = 0 // Insert the object
    17  	DeriveUpdate DeriveResult = 1 // Update the object (if it exists)
    18  	DeriveDelete DeriveResult = 2 // Delete the object
    19  	DeriveSkip   DeriveResult = 3 // Skip
    20  )
    21  
    22  type DeriveParams[In, Out any] struct {
    23  	cell.In
    24  
    25  	Lifecycle cell.Lifecycle
    26  	Jobs      job.Registry
    27  	Health    cell.Health
    28  	DB        *DB
    29  	InTable   Table[In]
    30  	OutTable  RWTable[Out]
    31  }
    32  
    33  // Derive constructs and registers a job to transform objects from the input table to the
    34  // output table, e.g. derive the output table from the input table. Useful when constructing
    35  // a reconciler that has its desired state solely derived from a single table. For example
    36  // the bandwidth manager's desired state is directly derived from the devices table.
    37  //
    38  // Derive is parametrized with the transform function that transforms the input object
    39  // into the output object. If the transform function returns false, then the object
    40  // is skipped.
    41  //
    42  // Example use:
    43  //
    44  //	cell.Invoke(
    45  //	  statedb.Derive[*tables.Device, *Foo](
    46  //	    func(d *Device, deleted bool) (*Foo, DeriveResult) {
    47  //	      if deleted {
    48  //	        return &Foo{Index: d.Index}, DeriveDelete
    49  //	      }
    50  //	      return &Foo{Index: d.Index}, DeriveInsert
    51  //	    }),
    52  //	)
    53  func Derive[In, Out any](jobName string, transform func(obj In, deleted bool) (Out, DeriveResult)) func(DeriveParams[In, Out]) {
    54  	return func(p DeriveParams[In, Out]) {
    55  		g := p.Jobs.NewGroup(p.Health)
    56  		g.Add(job.OneShot(
    57  			jobName,
    58  			derive[In, Out]{p, jobName, transform}.loop),
    59  		)
    60  		p.Lifecycle.Append(g)
    61  	}
    62  
    63  }
    64  
    65  type derive[In, Out any] struct {
    66  	DeriveParams[In, Out]
    67  	jobName   string
    68  	transform func(obj In, deleted bool) (Out, DeriveResult)
    69  }
    70  
    71  func (d derive[In, Out]) loop(ctx context.Context, _ cell.Health) error {
    72  	out := d.OutTable
    73  	txn := d.DB.WriteTxn(d.InTable)
    74  	iter, err := d.InTable.Changes(txn)
    75  	txn.Commit()
    76  	if err != nil {
    77  		return err
    78  	}
    79  	for {
    80  		wtxn := d.DB.WriteTxn(out)
    81  		changes, watch := iter.Next(wtxn)
    82  		for change := range changes {
    83  			outObj, result := d.transform(change.Object, change.Deleted)
    84  			switch result {
    85  			case DeriveInsert:
    86  				_, _, err = out.Insert(wtxn, outObj)
    87  			case DeriveUpdate:
    88  				_, _, found := out.Get(wtxn, out.PrimaryIndexer().QueryFromObject(outObj))
    89  				if found {
    90  					_, _, err = out.Insert(wtxn, outObj)
    91  				}
    92  			case DeriveDelete:
    93  				_, _, err = out.Delete(wtxn, outObj)
    94  			case DeriveSkip:
    95  			}
    96  			if err != nil {
    97  				wtxn.Abort()
    98  				return err
    99  			}
   100  		}
   101  		wtxn.Commit()
   102  
   103  		select {
   104  		case <-watch:
   105  		case <-ctx.Done():
   106  			return nil
   107  		}
   108  	}
   109  }