content: /*
   *
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you 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
   *
   *   http://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 org.apache.bookkeeper.common.util.affinity.impl;
  
  import java.io.Closeable;
  import java.io.File;
  import java.io.IOException;
  import java.io.RandomAccessFile;
  import java.nio.channels.FileChannel;
  import java.nio.channels.FileLock;
  import java.nio.file.Paths;
  import java.util.ArrayList;
  import java.util.List;
  import java.util.Set;
  import java.util.SortedSet;
  import java.util.TreeSet;
  import lombok.experimental.UtilityClass;
  import lombok.extern.slf4j.Slf4j;
  import org.apache.bookkeeper.common.util.nativelib.NativeUtils;
  
  /**
   * Implementation of CPU Affinity functionality.
   */
  @UtilityClass
  @Slf4j
  public class CpuAffinityImpl {
  
      private static boolean isInitialized = false;
      private static boolean isSupported;
  
      // Id of CPU cores acquired by this process
      private static final SortedSet<Integer> acquiredProcessors = new TreeSet<>();
  
      // Ids of processors that were isolated by Linux at boot time. This is the set
      // of processors that can acquired by this or other processes.
      private static SortedSet<Integer> isolatedProcessors = null;
  
      private static ProcessorsInfo processorsInfo = null;
  
      public static synchronized void acquireCore() {
          if (!isInitialized) {
              init();
          }
  
          if (!isSupported) {
              throw new RuntimeException("CPU Affinity not supported in current environment");
          }
  
          if (!CpuAffinityJni.isRoot()) {
              throw new RuntimeException("CPU Affinity can only be set if the process is running as root");
          }
  
          try {
              int cpu = pickAvailableCpu();
              CpuAffinityJni.setAffinity(cpu);
  
              log.info("Thread {} has successfully acquired ownership of cpu {}", Thread.currentThread().getName(), cpu);
          } catch (IOException e) {
              throw new RuntimeException("Failed to acquire CPU core: " + e.getMessage());
          }
      }
  
      private static final String LOCK_FILE_PREFIX = Paths.get(System.getProperty("java.io.tmpdir"), "cpu-lock-")
              .toString();
  
      /**
       * Other than the cores acquired by this process, there might be other processes on the same host trying to acquire
       * the available cores.
       *
       * <p>We use file-locks to ensure that other processes are aware of which CPUs are taken and that these locks are
       * automatically released if the process crashes.
       */
      private static synchronized int pickAvailableCpu() throws IOException {
          if (isolatedProcessors == null) {
              isolatedProcessors = IsolatedProcessors.get();
          }
          for (int isolatedCpu : isolatedProcessors) {
              if (log.isDebugEnabled()) {
                  log.debug("Checking CPU {}", isolatedCpu);
              }
              if (acquiredProcessors.contains(isolatedCpu)) {
                  if (log.isDebugEnabled()) {
                      log.debug("Ignoring CPU {} since it's already acquired", isolatedCpu);
                  }
                  continue;
              }
  
              if (tryAcquireCpu(isolatedCpu)) {
                  if (log.isDebugEnabled()) {
                      log.debug("Using CPU {}", isolatedCpu);
                  }
                  return isolatedCpu;
              }
          }
  
          throw new RuntimeException(
                  "There is no available isolated CPU to acquire for thread " + Thread.currentThread().getName());
      }
  
      private static boolean tryAcquireCpu(int targetCpu) throws IOException {
          // First, acquire lock on all the cpus that share the same core as target cpu
          if (processorsInfo == null) {
              processorsInfo = ProcessorsInfo.parseCpuInfo();
          }
  
          Set<Integer> cpusToAcquire = processorsInfo.getCpusOnSameCore(targetCpu);
          List<Closeable> acquiredCpus = new ArrayList<>();
  
          for (int cpu : cpusToAcquire) {
              Closeable lock = tryAcquireFileLock(cpu);
              if (lock == null) {
                  if (log.isDebugEnabled()) {
                      log.debug("Failed to acquire lock on CPU {}", cpu);
                  }
  
                  // Failed to acquire one cpu, release the rest that were already locked
                  for (Closeable l : acquiredCpus) {
                      l.close();
                  }
  
                  return false;
              } else {
                  acquiredCpus.add(lock);
              }
          }
  
          // At this point, we have ownership of all required cpus
          // Make sure the requested CPU is enabled and that all other CPUs on the same core are disabled, so that
          // hyper-threading will not affect latency.
          for (int cpu : cpusToAcquire) {
              if (cpu == targetCpu) {
                  IsolatedProcessors.enableCpu(cpu);
              } else {
                  IsolatedProcessors.disableCpu(cpu);
              }
  
              acquiredProcessors.add(cpu);
          }
          return true;
      }
  
      /**
       * Try to acquire a lock on a particular cpu.
       *
       * @return null if the lock was not available
       * @return a {@link Closeable} lock object if the lock was acquired
       */
      private static Closeable tryAcquireFileLock(int cpu) throws IOException {
          String lockPath = LOCK_FILE_PREFIX + cpu;
  
          RandomAccessFile file = null;
          FileChannel channel = null;
          FileLock lock = null;
  
          try {
              file = new RandomAccessFile(new File(lockPath), "rw");
              channel = file.getChannel();
              lock = channel.tryLock();
  
              if (lock == null) {
                  return null;
              } else {
                  final FileLock finalLock = lock;
                  final FileChannel finalChannel = channel;
                  final RandomAccessFile finalFile = file;
                  return () -> {
                      // Closable object
                      finalLock.close();
                      finalChannel.close();
                      finalFile.close();
                  };
              }
          } finally {
              if (lock == null) {
                  if (channel != null) {
                      channel.close();
                  }
  
                  if (file != null) {
                      file.close();
                  }
              }
          }
      }
  
      private static void init() {
          try {
              // Since this feature is only available in Linux, there's no point
              // in checking for MacOS jnilib or Windows dll extensions
              NativeUtils.loadLibraryFromJar("/lib/libcpu-affinity.so");
              isSupported = CpuAffinityJni.isAvailable();
          } catch (final Exception | UnsatisfiedLinkError e) {
              log.warn("Unable to load CPU affinity library: {}", e.getMessage(), e);
              isSupported = false;
          } finally {
              isInitialized = true;
          }
      }
  
  }
  