com.staros.shard.Shard Maven / Gradle / Ivy
// Copyright 2021-present StarRocks, Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package com.staros.shard;
import com.staros.filecache.FileCache;
import com.staros.filestore.FilePath;
import com.staros.proto.AddShardInfo;
import com.staros.proto.ReplicaInfo;
import com.staros.proto.ShardInfo;
import com.staros.proto.ShardState;
import com.staros.replica.Replica;
import com.staros.util.LockCloseable;
import com.staros.util.Text;
import com.staros.util.Writable;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.io.DataInput;
import java.io.DataOutput;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.locks.ReentrantLock;
import java.util.stream.Collectors;
public class Shard implements Writable {
private static final Logger LOG = LogManager.getLogger(Shard.class);
private final String serviceId;
// TODO: use Set to ensure uniqueness of groupIds
private List groupIds;
private final long shardId;
private ShardState state;
private FilePath filePath;
private FileCache fileCache;
private int expectedReplicaCount; // set when create shard
private List replicas; // actual replicas at current time
private Map properties;
private final ReentrantLock lock;
public Shard(String serviceId,
List groupIds,
long shardId) {
this(serviceId, groupIds, shardId, null, null);
}
public Shard(
String serviceId,
List groupIds,
long shardId,
FilePath filePath,
FileCache fileCache) {
this.serviceId = serviceId;
this.groupIds = Collections.unmodifiableList(new ArrayList<>(groupIds));
this.shardId = shardId;
this.state = ShardState.NORMAL;
this.filePath = filePath;
this.fileCache = fileCache;
this.expectedReplicaCount = 1;
this.replicas = new ArrayList<>();
this.properties = new HashMap<>();
this.lock = new ReentrantLock();
}
public String getServiceId() {
return serviceId;
}
public List getGroupIds() {
try (LockCloseable ignore = new LockCloseable(lock)) {
// It is safe to return a read-only list to outside.
return groupIds;
}
}
public long getShardId() {
return shardId;
}
public boolean joinGroup(long groupId) {
try (LockCloseable ignore = new LockCloseable(lock)) {
if (groupIds.contains(groupId)) {
return false;
}
// Copy-On-Write
List newGroup = new ArrayList<>(groupIds);
newGroup.add(groupId);
groupIds = Collections.unmodifiableList(newGroup);
return true;
}
}
public boolean quitGroup(long groupId) {
try (LockCloseable ignore = new LockCloseable(lock)) {
if (!groupIds.contains(groupId)) {
return false;
}
//Copy-On-Write
List newGroup = new ArrayList<>(groupIds);
newGroup.remove(groupId);
groupIds = Collections.unmodifiableList(newGroup);
return true;
}
}
public ShardState getState() {
return state;
}
public void setState(ShardState state) {
this.state = state;
}
public FilePath getFilePath() {
return filePath;
}
public FileCache getFileCache() {
return fileCache;
}
public void setReplicas(List replicas) {
this.replicas = replicas;
}
public List getReplicaWorkerIds() {
try (LockCloseable ignore = new LockCloseable(lock)) {
return replicas.stream().map(Replica::getWorkerId).collect(Collectors.toList());
}
}
public int getReplicaSize() {
try (LockCloseable ignore = new LockCloseable(lock)) {
return replicas.size();
}
}
public int getExpectedReplicaCount() {
return expectedReplicaCount;
}
public void setExpectedReplicaCount(int expectedReplicaCount) {
this.expectedReplicaCount = expectedReplicaCount;
}
public boolean hasReplica(long workerId) {
try (LockCloseable ignore = new LockCloseable(lock)) {
return replicas.stream().anyMatch(x -> x.getWorkerId() == workerId);
}
}
public boolean addReplica(long workerId) {
try (LockCloseable ignore = new LockCloseable(lock)) {
for (Replica replica : replicas) {
if (replica.getWorkerId() == workerId) {
return false;
}
}
// TODO: every Replica is PRIMARY for now as the role is not used yet.
// Proper PRIMARY/SECONDARY should take worker group into consideration.
replicas.add(new Replica(workerId));
return true;
}
}
public boolean removeReplica(long workerId) {
try (LockCloseable ignore = new LockCloseable(lock)) {
// TODO: properly set replica ROLE when multiple worker groups is taken into consideration.
return replicas.removeIf(x -> x.getWorkerId() == workerId);
}
}
public void setProperties(Map properties) {
this.properties = properties;
}
public Map getProperties() {
return properties;
}
public ShardInfo toProtobuf() {
try (LockCloseable ignore = new LockCloseable(lock)) {
ShardInfo.Builder builder = ShardInfo.newBuilder();
builder.setServiceId(serviceId);
builder.addAllGroupIds(groupIds);
builder.setShardId(shardId);
builder.setShardState(state);
if (filePath != null) {
builder.setFilePath(filePath.toProtobuf());
}
if (fileCache != null) {
builder.setFileCache(fileCache.toProtobuf());
}
for (Replica replica : replicas) {
builder.addReplicaInfo(replica.toProtobuf());
}
builder.putAllShardProperties(properties);
return builder.build();
}
}
public AddShardInfo getAddShardInfo() {
try (LockCloseable ignore = new LockCloseable(lock)) {
return AddShardInfo.newBuilder()
.setShardId(shardId)
.setFileCacheInfo(fileCache.toProtobuf())
.setFilePathInfo((filePath.toProtobuf()))
.putAllShardProperties(properties)
.build();
}
}
public static Shard fromProtobuf(ShardInfo info) {
String serviceId = info.getServiceId();
List groupIds = info.getGroupIdsList();
long shardId = info.getShardId();
FilePath path = FilePath.fromProtobuf(info.getFilePath());
FileCache cache = FileCache.fromProtobuf(info.getFileCache());
Shard shard = new Shard(serviceId, groupIds, shardId, path, cache);
shard.setState(info.getShardState());
List replicaInfos = info.getReplicaInfoList();
List replicas = new ArrayList<>(replicaInfos.size());
for (ReplicaInfo replicaInfo : replicaInfos) {
replicas.add(Replica.fromProtobuf(replicaInfo));
}
shard.setReplicas(replicas);
Map properties = info.getShardPropertiesMap();
shard.setProperties(properties);
return shard;
}
@Override
public void write(DataOutput out) throws IOException {
byte[] bytes = toProtobuf().toByteArray();
Text.writeBytes(out, bytes);
}
public static Shard read(DataInput in) throws IOException {
byte[] bytes = Text.readBytes(in);
ShardInfo info = ShardInfo.parseFrom(bytes);
return Shard.fromProtobuf(info);
}
}
© 2015 - 2025 Weber Informatics LLC | Privacy Policy