github.com/whtcorpsinc/milevadb-prod@v0.0.0-20211104133533-f57f4be3b597/causetstore/petri/acyclic/tenant/mock.go (about)

     1  // Copyright 2020 WHTCORPS INC, 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  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package tenant
    15  
    16  import (
    17  	"context"
    18  	"sync/atomic"
    19  
    20  	"github.com/whtcorpsinc/errors"
    21  )
    22  
    23  var _ Manager = &mockManager{}
    24  
    25  // mockManager represents the structure which is used for electing tenant.
    26  // It's used for local causetstore and testing.
    27  // So this worker will always be the tenant.
    28  type mockManager struct {
    29  	tenant  int32
    30  	id     string // id is the ID of manager.
    31  	cancel context.CancelFunc
    32  }
    33  
    34  // NewMockManager creates a new mock Manager.
    35  func NewMockManager(ctx context.Context, id string) Manager {
    36  	_, cancelFunc := context.WithCancel(ctx)
    37  	return &mockManager{
    38  		id:     id,
    39  		cancel: cancelFunc,
    40  	}
    41  }
    42  
    43  // ID implements Manager.ID interface.
    44  func (m *mockManager) ID() string {
    45  	return m.id
    46  }
    47  
    48  // IsTenant implements Manager.IsTenant interface.
    49  func (m *mockManager) IsTenant() bool {
    50  	return atomic.LoadInt32(&m.tenant) == 1
    51  }
    52  
    53  func (m *mockManager) toBeTenant() {
    54  	atomic.StoreInt32(&m.tenant, 1)
    55  }
    56  
    57  // RetireTenant implements Manager.RetireTenant interface.
    58  func (m *mockManager) RetireTenant() {
    59  	atomic.StoreInt32(&m.tenant, 0)
    60  }
    61  
    62  // Cancel implements Manager.Cancel interface.
    63  func (m *mockManager) Cancel() {
    64  	m.cancel()
    65  }
    66  
    67  // GetTenantID implements Manager.GetTenantID interface.
    68  func (m *mockManager) GetTenantID(ctx context.Context) (string, error) {
    69  	if m.IsTenant() {
    70  		return m.ID(), nil
    71  	}
    72  	return "", errors.New("no tenant")
    73  }
    74  
    75  // CampaignTenant implements Manager.CampaignTenant interface.
    76  func (m *mockManager) CampaignTenant() error {
    77  	m.toBeTenant()
    78  	return nil
    79  }
    80  
    81  // ResignTenant lets the tenant start a new election.
    82  func (m *mockManager) ResignTenant(ctx context.Context) error {
    83  	if m.IsTenant() {
    84  		m.RetireTenant()
    85  	}
    86  	return nil
    87  }